Handle None in assert_valid_xy (#6871)

* Handle none in assert_valid_xy

* add test

* Update xarray/tests/test_plot.py

* Update test_plot.py

Co-authored-by: Deepak Cherian <[email protected]>

Author: Illviljan

xarray/plot/utils.py

@@ -5,7 +5,7 @@
 import warnings
 from datetime import datetime
 from inspect import getfullargspec
-from typing import Any, Iterable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Hashable, Iterable, Mapping, Sequence
 
 import numpy as np
 import pandas as pd
@@ -28,6 +28,11 @@
 except ImportError:
     cftime = None
 
+
+if TYPE_CHECKING:
+    from ..core.dataarray import DataArray
+
+
 ROBUST_PERCENTILE = 2.0
 
 # copied from seaborn
@@ -396,7 +401,8 @@ def _infer_xy_labels(darray, x, y, imshow=False, rgb=None):
     return x, y
 
 
-def _assert_valid_xy(darray, xy, name):
+# TODO: Can by used to more than x or y, rename?
+def _assert_valid_xy(darray: DataArray, xy: None | Hashable, name: str) -> None:
     """
     make sure x and y passed to plotting functions are valid
     """
@@ -410,9 +416,11 @@ def _assert_valid_xy(darray, xy, name):
 
     valid_xy = (set(darray.dims) | set(darray.coords)) - multiindex_dims
 
-    if xy not in valid_xy:
-        valid_xy_str = "', '".join(sorted(valid_xy))
-        raise ValueError(f"{name} must be one of None, '{valid_xy_str}'")
+    if (xy is not None) and (xy not in valid_xy):
+        valid_xy_str = "', '".join(sorted(tuple(str(v) for v in valid_xy)))
+        raise ValueError(
+            f"{name} must be one of None, '{valid_xy_str}'. Received '{xy}' instead."
+        )
 
 
 def get_axis(figsize=None, size=None, aspect=None, ax=None, **kwargs):
@@ -1152,7 +1160,7 @@ def _adjust_legend_subtitles(legend):
 
 def _infer_meta_data(ds, x, y, hue, hue_style, add_guide, funcname):
     dvars = set(ds.variables.keys())
-    error_msg = f" must be one of ({', '.join(dvars)})"
+    error_msg = f" must be one of ({', '.join(sorted(tuple(str(v) for v in dvars)))})"
 
     if x not in dvars:
         raise ValueError(f"Expected 'x' {error_msg}. Received {x} instead.")

xarray/tests/test_plot.py

@@ -17,6 +17,7 @@
 from xarray.plot.dataset_plot import _infer_meta_data
 from xarray.plot.plot import _infer_interval_breaks
 from xarray.plot.utils import (
+    _assert_valid_xy,
     _build_discrete_cmap,
     _color_palette,
     _determine_cmap_params,
@@ -3025,3 +3026,19 @@ def test_datarray_scatter(x, y, z, hue, markersize, row, col, add_legend, add_co
             add_legend=add_legend,
             add_colorbar=add_colorbar,
         )
+
+
+@requires_matplotlib
+def test_assert_valid_xy() -> None:
+    ds = xr.tutorial.scatter_example_dataset()
+    darray = ds.A
+
+    # x is valid and should not error:
+    _assert_valid_xy(darray=darray, xy="x", name="x")
+
+    # None should be valid as well even though it isn't in the valid list:
+    _assert_valid_xy(darray=darray, xy=None, name="x")
+
+    # A hashable that is not valid should error:
+    with pytest.raises(ValueError, match="x must be one of"):
+        _assert_valid_xy(darray=darray, xy="error_now", name="x")