Adjust add_colorbar option for facets

Author: shoyer

xray/plot/facetgrid.py

@@ -221,14 +221,10 @@ def map_dataarray(self, func, x, y, **kwargs):
 
         cmap_params = _determine_cmap_params(**cmap_kwargs)
 
-        defaults = {
-            'add_colorbar': False,
-            'add_labels': False,
-        }
-
         # Order is important
-        defaults.update(kwargs)
-        defaults.update(cmap_params)
+        func_kwargs = kwargs.copy()
+        func_kwargs.update(cmap_params)
+        func_kwargs.update({'add_colorbar': False, 'add_labels': False})
 
         # Get x, y labels for the first subplot
         x, y = _infer_xy_labels(darray=self.data.loc[self.name_dicts.flat[0]],
@@ -238,13 +234,12 @@ def map_dataarray(self, func, x, y, **kwargs):
             # None is the sentinel value
             if d is not None:
                 subset = self.data.loc[d]
-                mappable = func(subset, x, y, ax=ax, **defaults)
+                mappable = func(subset, x, y, ax=ax, **func_kwargs)
                 self._mappables.append(mappable)
 
-        self._cmap_extend = defaults.get('extend')
+        self._cmap_extend = cmap_params.get('extend')
         self._finalize_grid(x, y)
 
-        # colorbar
         if kwargs.get('add_colorbar', True):
             self.add_colorbar()
 
@@ -266,12 +261,10 @@ def add_colorbar(self, **kwargs):
         kwargs = kwargs.copy()
         if self._cmap_extend is not None:
             kwargs.setdefault('extend', self._cmap_extend)
-        cbar = self.fig.colorbar(self._mappables[-1],
-                                 ax=list(self.axes.flat),
-                                 **kwargs)
-        if getattr(self.data, 'name', False):
-            cbar.set_label(self.data.name, rotation=90,
-                           verticalalignment='bottom')
+        if getattr(self.data, 'name', None) is not None:
+            kwargs.setdefault('label', self.data.name)
+        self.fig.colorbar(self._mappables[-1], ax=list(self.axes.flat),
+                          **kwargs)
         return self
 
     def set_axis_labels(self, x_var=None, y_var=None):

xray/plot/plot.py

@@ -343,10 +343,6 @@ def newplotfunc(darray, x=None, y=None, ax=None, row=None, col=None,
             allargs = locals().copy()
             allargs.update(allargs.pop('kwargs'))
 
-            # Allows use of better FacetGrid defaults
-            assert allargs.pop('add_labels')
-            assert allargs.pop('add_colorbar')
-
             # Need the decorated plotting function
             allargs['plotfunc'] = globals()[plotfunc.__name__]
 

xray/test/test_plot.py

@@ -32,7 +32,8 @@ def text_in_fig():
     return set(alltxt)
 
 
-def find_colorbars():
+def find_possible_colorbars():
+    # nb. this function also matches meshes from pcolormesh
     return plt.gcf().findobj(mpl.collections.QuadMesh)
 
 
@@ -794,7 +795,7 @@ def test_colorbar(self):
             clim = np.array(image.get_clim())
             self.assertTrue(np.allclose(expected, clim))
 
-        self.assertEqual(1, len(find_colorbars()))
+        self.assertEqual(1, len(find_possible_colorbars()))
 
     def test_empty_cell(self):
         g = xplt.FacetGrid(self.darray, col='z', col_wrap=2)
@@ -899,11 +900,11 @@ def test_map_dataset(self):
 
         # colorbar can't be inferred automatically
         self.assertNotIn('foo', alltxt)
-        self.assertEqual(0, len(find_colorbars()))
+        self.assertEqual(0, len(find_possible_colorbars()))
 
         g.add_colorbar(label='colors!')
         self.assertIn('colors!', text_in_fig())
-        self.assertEqual(1, len(find_colorbars()))
+        self.assertEqual(1, len(find_possible_colorbars()))
 
     def test_set_axis_labels(self):
         g = self.g.map_dataarray(xplt.contourf, 'x', 'y')
@@ -912,6 +913,19 @@ def test_set_axis_labels(self):
         for label in ['longitude', 'latitude']:
             self.assertIn(label, alltxt)
 
+    def test_facetgrid_colorbar(self):
+        a = easy_array((10, 15, 4))
+        d = DataArray(a, dims=['y', 'x', 'z'], name='foo')
+
+        d.plot.imshow(x='x', y='y', col='z')
+        self.assertEqual(1, len(find_possible_colorbars()))
+
+        d.plot.imshow(x='x', y='y', col='z', add_colorbar=True)
+        self.assertEqual(1, len(find_possible_colorbars()))
+
+        d.plot.imshow(x='x', y='y', col='z', add_colorbar=False)
+        self.assertEqual(0, len(find_possible_colorbars()))
+
 
 class TestFacetGrid4d(PlotTestCase):