Rework plotting heuristics. (#251)

Author: dcherian

cf_xarray/accessor.py

@@ -736,7 +736,7 @@ def check_results(names, key):
         )
 
 
-def _possible_x_y_plot(obj, key):
+def _possible_x_y_plot(obj, key, skip=None):
     """Guesses a name for an x/y variable if possible."""
     # in priority order
     x_criteria = [
@@ -759,11 +759,20 @@ def _get_possible(accessor, criteria):
         from xarray.core.utils import is_scalar
 
         for attr, key in criteria:
-            value = getattr(accessor, attr).get(key)
-            if not value or len(value) > 1:
+            values = getattr(accessor, attr).get(key)
+            ax_coord_name = getattr(accessor, attr).get(key)
+            if not values:
                 continue
-            if not is_scalar(accessor._obj[value[0]]):
-                return value[0]
+            elif ax_coord_name:
+                values = [v for v in values if v in ax_coord_name]
+
+            values = [v for v in values if v != skip]
+            if len(values) == 1 and not is_scalar(accessor._obj[values[0]]):
+                return values[0]
+            else:
+                for v in values:
+                    if not is_scalar(accessor._obj[v]):
+                        return v
         return None
 
     if key == "x":
@@ -825,9 +834,9 @@ def _plot_decorator(self, func):
 
         @functools.wraps(func)
         def _plot_wrapper(*args, **kwargs):
-            def _process_x_or_y(kwargs, key):
+            def _process_x_or_y(kwargs, key, skip=None):
                 if key not in kwargs:
-                    kwargs[key] = _possible_x_y_plot(self._obj, key)
+                    kwargs[key] = _possible_x_y_plot(self._obj, key, skip)
 
                 value = kwargs.get(key)
                 if value:
@@ -847,13 +856,15 @@ def _process_x_or_y(kwargs, key):
                 and (kwargs.get("hue") or self._obj.ndim == 1)
             )
             if is_line_plot:
-                if not kwargs.get("hue"):
-                    kwargs = _process_x_or_y(kwargs, "x")
+                hue = kwargs.get("hue")
+                if "x" not in kwargs and "y" not in kwargs:
+                    kwargs = _process_x_or_y(kwargs, "x", skip=hue)
                     if not kwargs.get("x"):
-                        kwargs = _process_x_or_y(kwargs, "y")
+                        kwargs = _process_x_or_y(kwargs, "y", skip=hue)
+
             else:
-                kwargs = _process_x_or_y(kwargs, "x")
-                kwargs = _process_x_or_y(kwargs, "y")
+                kwargs = _process_x_or_y(kwargs, "x", skip=kwargs.get("y"))
+                kwargs = _process_x_or_y(kwargs, "y", skip=kwargs.get("x"))
 
             return func(*args, **kwargs)
 

cf_xarray/tests/test_accessor.py

@@ -496,9 +496,9 @@ def test_dataarray_plot():
     np.testing.assert_equal(rv[0].get_xdata(), obj.lon.data)
     plt.close()
 
-    # rv = obj.cf.isel(T=1, Y=[0, 1, 2]).cf.plot(hue="Y")
-    # np.testing.assert_equal(rv[0].get_xdata(), obj.lon.data)
-    # plt.close()
+    rv = obj.cf.isel(T=1, Y=[0, 1, 2]).cf.plot(hue="Y")
+    np.testing.assert_equal(rv[0].get_xdata(), obj.lon.data)
+    plt.close()
 
     rv = obj.cf.isel(T=1, Y=[0, 1, 2]).cf.plot.line()
     np.testing.assert_equal(rv[0].get_xdata(), obj.lon.data)