Test regexes for _guess_coord_axis (#172)

Author: dcherian

cf_xarray/accessor.py

@@ -116,7 +116,7 @@
 
 #: regular expressions for guess_coord_axis
 regex = {
-    "time": "time[0-9]*|min|hour|day|week|month|year",
+    "time": "(time|min|hour|day|week|month|year)[0-9]*",
     "vertical": (
         "(lv_|bottom_top|sigma|h(ei)?ght|altitude|depth|isobaric|pres|"
         "isotherm)[a-z_]*[0-9]*"
@@ -130,16 +130,16 @@
 regex["T"] = regex["time"]
 
 
-attrs = {
+ATTRS = {
     "X": {"axis": "X"},
     "T": {"axis": "T", "standard_name": "time"},
     "Y": {"axis": "Y"},
     "Z": {"axis": "Z"},
     "latitude": {"units": "degrees_north", "standard_name": "latitude"},
     "longitude": {"units": "degrees_east", "standard_name": "longitude"},
 }
-attrs["time"] = attrs["T"]
-attrs["vertical"] = attrs["Z"]
+ATTRS["time"] = ATTRS["T"]
+ATTRS["vertical"] = ATTRS["Z"]
 
 
 # Type for Mapper functions
@@ -1361,7 +1361,7 @@ def guess_coord_axis(self, verbose: bool = False) -> Union[DataArray, Dataset]:
                     print(
                         f"I think {var!r} is of type 'time'. It has a datetime-like type."
                     )
-                obj[var].attrs = dict(ChainMap(obj[var].attrs, attrs["time"]))
+                obj[var].attrs = dict(ChainMap(obj[var].attrs, ATTRS["time"]))
                 continue  # prevent second detection
 
             for axis, pattern in regex.items():
@@ -1371,7 +1371,7 @@ def guess_coord_axis(self, verbose: bool = False) -> Union[DataArray, Dataset]:
                         print(
                             f"I think {var!r} is of type {axis!r}. It matched {pattern!r}"
                         )
-                    obj[var].attrs = dict(ChainMap(obj[var].attrs, attrs[axis]))
+                    obj[var].attrs = dict(ChainMap(obj[var].attrs, ATTRS[axis]))
         return obj
 
 

cf_xarray/tests/test_accessor.py

@@ -1,3 +1,4 @@
+import itertools
 from textwrap import dedent
 
 import matplotlib as mpl
@@ -627,26 +628,77 @@ def test_docstring():
     assert "One or more of ('X'" in airds.cf.mean.__doc__
 
 
-def test_guess_coord_axis():
+def _make_names(prefixes):
+    suffixes = ["", "a", "_a", "0", "_0"]
+    return [
+        f"{prefix}{suffix}" for prefix, suffix in itertools.product(prefixes, suffixes)
+    ]
+
+
+_TIME_NAMES = _make_names(
+    [
+        "time",
+        "min",
+        "hour",
+        "day",
+        "week",
+        "month",
+        "year",
+    ]
+)
+_VERTICAL_NAMES = _make_names(
+    [
+        "lv_1",
+        "bottom_top",
+        "sigma",
+        "sigma_w",
+        "hght",
+        "height",
+        "altitude",
+        "depth",
+        "isobaric",
+        "pressure",
+        "isotherm",
+    ]
+)
+_X_NAMES = _make_names(["x"])
+_Y_NAMES = _make_names(["y"])
+_Z_NAMES = _VERTICAL_NAMES
+_LATITUDE_NAMES = _make_names(["lat", "latitude"])
+_LONGITUDE_NAMES = _make_names(["lon", "longitude"])
+
+
+@pytest.mark.parametrize(
+    "kind, names",
+    [
+        ["X", _X_NAMES],
+        ["Y", _Y_NAMES],
+        ["Z", _Z_NAMES],
+        ["T", _TIME_NAMES],
+        ["latitude", _LATITUDE_NAMES],
+        ["longitude", _LONGITUDE_NAMES],
+    ],
+)
+def test_guess_coord_axis(kind, names):
+    from cf_xarray.accessor import ATTRS
+
+    for varname in names:
+        ds = xr.Dataset()
+        ds[varname] = (varname, [1, 2, 3, 4, 5])
+        dsnew = ds.cf.guess_coord_axis()
+        assert dsnew[varname].attrs == ATTRS[kind]
+
+        varname = varname.upper()
+        ds[varname] = (varname, [1, 2, 3, 4, 5])
+        dsnew = ds.cf.guess_coord_axis()
+        assert dsnew[varname].attrs == ATTRS[kind]
+
+
+def test_guess_coord_axis_datetime():
     ds = xr.Dataset()
     ds["time"] = ("time", pd.date_range("2001-01-01", "2001-04-01"))
-    ds["lon_rho"] = ("lon_rho", [1, 2, 3, 4, 5])
-    ds["lat_rho"] = ("lat_rho", [1, 2, 3, 4, 5])
-    ds["x1"] = ("x1", [1, 2, 3, 4, 5])
-    ds["y1"] = ("y1", [1, 2, 3, 4, 5])
-
     dsnew = ds.cf.guess_coord_axis()
     assert dsnew.time.attrs == {"standard_name": "time", "axis": "T"}
-    assert dsnew.lon_rho.attrs == {
-        "standard_name": "longitude",
-        "units": "degrees_east",
-    }
-    assert dsnew.lat_rho.attrs == {
-        "standard_name": "latitude",
-        "units": "degrees_north",
-    }
-    assert dsnew.x1.attrs == {"axis": "X"}
-    assert dsnew.y1.attrs == {"axis": "Y"}
 
 
 def test_attributes():