Bugfix for faceting line plots. (#2229)

* Bugfix for faceting line plots.

* Make add_legend public and usable.

Author: dcherian

xarray/plot/facetgrid.py

@@ -293,14 +293,16 @@ def map_dataarray_line(self, x=None, y=None, hue=None, **kwargs):
                                 ax=ax, _labels=False,
                                 **kwargs)
                 self._mappables.append(mappable)
-        _, _, _, xlabel, ylabel, huelabel = _infer_line_data(
+        _, _, hueplt, xlabel, ylabel, huelabel = _infer_line_data(
             darray=self.data.loc[self.name_dicts.flat[0]],
             x=x, y=y, hue=hue)
 
+        self._hue_var = hueplt
+        self._hue_label = huelabel
         self._finalize_grid(xlabel, ylabel)
 
-        if add_legend and huelabel:
-            self.add_line_legend(huelabel)
+        if add_legend and hueplt is not None and huelabel is not None:
+            self.add_legend()
 
         return self
 
@@ -314,12 +316,12 @@ def _finalize_grid(self, *axlabels):
             if namedict is None:
                 ax.set_visible(False)
 
-    def add_line_legend(self, huelabel):
+    def add_legend(self, **kwargs):
         figlegend = self.fig.legend(
             handles=self._mappables[-1],
-            labels=list(self.data.coords[huelabel].values),
-            title=huelabel,
-            loc="center right")
+            labels=list(self._hue_var.values),
+            title=self._hue_label,
+            loc="center right", **kwargs)
 
         # Draw the plot to set the bounding boxes correctly
         self.fig.draw(self.fig.canvas.get_renderer())

xarray/tests/test_plot.py

@@ -1529,6 +1529,12 @@ def setUp(self):
                                         range(3), ['A', 'B', 'C', 'C++']],
                                 name='Cornelius Ortega the 1st')
 
+        self.darray.hue.name = 'huename'
+        self.darray.hue.attrs['units'] = 'hunits'
+        self.darray.x.attrs['units'] = 'xunits'
+        self.darray.col.attrs['units'] = 'colunits'
+        self.darray.row.attrs['units'] = 'rowunits'
+
     def test_facetgrid_shape(self):
         g = self.darray.plot(row='row', col='col', hue='hue')
         assert g.axes.shape == (len(self.darray.row), len(self.darray.col))