add functionality to declarative plotting

metpy/plots/declarative.py

@@ -673,9 +673,12 @@ class Plot2D(HasTraits):
     level = Union([Int(allow_none=True, default_value=None), Instance(units.Quantity)])
     time = Instance(datetime, allow_none=True)
 
+    contours = Union([List(Float()), Int()], default_value=25)
+    clabels = Bool(default_value=False)
     colormap = Unicode(allow_none=True, default_value=None)
     image_range = Union([Tuple(Int(allow_none=True), Int(allow_none=True)),
                          Instance(plt.Normalize)], default_value=(None, None))
+    colorbar = Unicode(default_value=None, allow_none=True)
 
     @property
     def _cmap_obj(self):
@@ -705,14 +708,23 @@ def clear(self):
 
         """
         if getattr(self, 'handle', None) is not None:
-            self.clear_handle()
+            if getattr(self.handle, 'collections', None) is not None:
+                self.clear_collections()
+            else:
+                self.clear_handle()
             self._need_redraw = True
 
     def clear_handle(self):
         """Clear the handle to the plot instance."""
         self.handle.remove()
         self.handle = None
 
+    def clear_collections(self):
+        """Clear the handle collections to the plot instance."""
+        for col in self.handle.collections:
+            col.remove()
+        self.handle = None
+
     @observe('parent')
     def _parent_changed(self, _):
         """Handle setting the parent object for the plot."""
@@ -795,6 +807,9 @@ def draw(self):
         if self._need_redraw:
             if getattr(self, 'handle', None) is None:
                 self._build()
+            if self.colorbar is not None:
+                self.parent.ax.figure.colorbar(
+                    self.handle, orientation=self.colorbar, pad=0, aspect=50)
             self._need_redraw = False
 
 
@@ -810,6 +825,13 @@ def _set_need_redraw(self, _):
             self.handle.set_norm(self._norm_obj)
             self._need_redraw = True
 
+    @observe('colorbar')
+    def _set_need_rebuild(self, _):
+        """Handle changes to attributes that need to regenerate everything."""
+        # Because matplotlib doesn't let you just change these properties, we need
+        # to trigger a clear and re-call of contour()
+        self.clear()
+
     @property
     def plotdata(self):
         """Return the data for plotting.
@@ -846,24 +868,41 @@ class ContourPlot(Plot2D):
 
     linecolor = Unicode('black')
     linewidth = Int(2)
-    contours = Union([List(Float()), Int()], default_value=25)
+    linestyle = Unicode('solid', allow_none=True)
 
-    @observe('contours', 'linecolor', 'linewidth')
+    @observe('contours', 'linecolor', 'linewidth', 'linestyle', 'clabels')
     def _set_need_rebuild(self, _):
         """Handle changes to attributes that need to regenerate everything."""
         # Because matplotlib doesn't let you just change these properties, we need
         # to trigger a clear and re-call of contour()
         self.clear()
 
-    def clear_handle(self):
-        """Clear the handle to the plot instance."""
-        for col in self.handle.collections:
-            col.remove()
-        self.handle = None
-
     def _build(self):
         """Build the plot by calling any plotting methods as necessary."""
         x, y, imdata = self.plotdata
         self.handle = self.parent.ax.contour(x, y, imdata, self.contours,
                                              colors=self.linecolor, linewidths=self.linewidth,
+                                             linestyles=self.linestyle,
                                              transform=imdata.metpy.cartopy_crs)
+        if self.clabels:
+            self.handle.clabel(inline=1, fmt='%.0f', inline_spacing=2,
+                               use_clabeltext=True)
+
+
+@exporter.export
+class FilledContourPlot(Plot2D):
+    """Represent a contour fill plot."""
+
+    @observe('contours', 'colorbar', 'colormap')
+    def _set_need_rebuild(self, _):
+        """Handle changes to attributes that need to regenerate everything."""
+        # Because matplotlib doesn't let you just change these properties, we need
+        # to trigger a clear and re-call of contour()
+        self.clear()
+
+    def _build(self):
+        """Build the plot by calling any plotting methods as necessary."""
+        x, y, imdata = self.plotdata
+        self.handle = self.parent.ax.contourf(x, y, imdata, self.contours,
+                                              cmap=self._cmap_obj, norm=self._norm_obj,
+                                              transform=imdata.metpy.cartopy_crs)

metpy/plots/tests/test_declarative.py

@@ -15,7 +15,7 @@
 
 from metpy.cbook import get_test_data
 from metpy.io import GiniFile
-from metpy.plots import ContourPlot, ImagePlot, MapPanel, PanelContainer
+from metpy.plots import ContourPlot, FilledContourPlot, ImagePlot, MapPanel, PanelContainer
 # Fixtures to make sure we have the right backend
 from metpy.testing import set_agg_backend  # noqa: F401, I202
 from metpy.units import units
@@ -74,6 +74,35 @@ def test_declarative_contour():
     return pc.figure
 
 
+@pytest.mark.mpl_image_compare(remove_text=True, tolerance=0.035)
+def test_declarative_contour_options():
+    """Test making a contour plot."""
+    data = xr.open_dataset(get_test_data('narr_example.nc', as_file_obj=False))
+
+    contour = ContourPlot()
+    contour.data = data
+    contour.field = 'Temperature'
+    contour.level = 700 * units.hPa
+    contour.contours = 30
+    contour.linewidth = 1
+    contour.linecolor = 'red'
+    contour.linestyle = 'dashed'
+    contour.clabels = True
+
+    panel = MapPanel()
+    panel.area = 'us'
+    panel.proj = 'lcc'
+    panel.layers = ['coastline', 'borders', 'usstates']
+    panel.plots = [contour]
+
+    pc = PanelContainer()
+    pc.size = (8, 8)
+    pc.panels = [panel]
+    pc.draw()
+
+    return pc.figure
+
+
 @pytest.mark.mpl_image_compare(remove_text=True, tolerance=0.016)
 def test_declarative_events():
     """Test that resetting traitlets properly propagates."""
@@ -112,6 +141,7 @@ def test_declarative_events():
     contour.field = 'Specific_humidity'
     img.field = 'Geopotential_height'
     img.colormap = 'plasma'
+    img.colorbar = 'horizontal'
 
     return pc.figure
 
@@ -150,6 +180,81 @@ def test_projection_object():
     return pc.figure
 
 
+@pytest.mark.mpl_image_compare(remove_text=True, tolerance=0.016)
+def test_colorfill():
+    """Test that we can use ContourFillPlot."""
+    data = xr.open_dataset(get_test_data('narr_example.nc', as_file_obj=False))
+
+    contour = FilledContourPlot()
+    contour.data = data
+    contour.level = 700 * units.hPa
+    contour.field = 'Temperature'
+    contour.colormap = 'coolwarm'
+    contour.colorbar = 'vertical'
+
+    panel = MapPanel()
+    panel.area = (-110, -60, 25, 55)
+    panel.layers = [cfeature.STATES]
+    panel.plots = [contour]
+
+    pc = PanelContainer()
+    pc.panel = panel
+    pc.size = (12, 8)
+    pc.draw()
+
+    return pc.figure
+
+
+@pytest.mark.mpl_image_compare(remove_text=True, tolerance=0.016)
+def test_colorfill_horiz_colorbar():
+    """Test that we can use ContourFillPlot."""
+    data = xr.open_dataset(get_test_data('narr_example.nc', as_file_obj=False))
+
+    contour = FilledContourPlot()
+    contour.data = data
+    contour.level = 700 * units.hPa
+    contour.field = 'Temperature'
+    contour.colormap = 'coolwarm'
+    contour.colorbar = 'horizontal'
+
+    panel = MapPanel()
+    panel.area = (-110, -60, 25, 55)
+    panel.layers = [cfeature.STATES]
+    panel.plots = [contour]
+
+    pc = PanelContainer()
+    pc.panel = panel
+    pc.size = (8, 8)
+    pc.draw()
+
+    return pc.figure
+
+
+@pytest.mark.mpl_image_compare(remove_text=True, tolerance=0.016)
+def test_colorfill_no_colorbar():
+    """Test that we can use ContourFillPlot."""
+    data = xr.open_dataset(get_test_data('narr_example.nc', as_file_obj=False))
+
+    contour = FilledContourPlot()
+    contour.data = data
+    contour.level = 700 * units.hPa
+    contour.field = 'Temperature'
+    contour.colormap = 'coolwarm'
+    contour.colorbar = None
+
+    panel = MapPanel()
+    panel.area = (-110, -60, 25, 55)
+    panel.layers = [cfeature.STATES]
+    panel.plots = [contour]
+
+    pc = PanelContainer()
+    pc.panel = panel
+    pc.size = (8, 8)
+    pc.draw()
+
+    return pc.figure
+
+
 @pytest.mark.mpl_image_compare(remove_text=True, tolerance=1.23)
 def test_global():
     """Test that we can set global extent."""
@@ -158,6 +263,7 @@ def test_global():
     img = ImagePlot()
     img.data = data
     img.field = 'IR'
+    img.colorbar = None
 
     panel = MapPanel()
     panel.area = 'global'
@@ -181,6 +287,7 @@ def test_latlon():
     img.field = 'Temperature_isobaric'
     img.level = 500 * units.hPa
     img.time = datetime(2017, 9, 5, 15, 0, 0)
+    img.colorbar = None
 
     contour = ContourPlot()
     contour.data = data