update

pcmdi_metrics/io/__init__.py

@@ -4,7 +4,7 @@
 from .base import MV2Json  # noqa
 from .default_regions_define import load_regions_specs  # noqa
 from .default_regions_define import region_subset  # noqa
-from .xcdat_xarray_dataset_io import (  # noqa
+from .xcdat_dataset_io import (  # noqa
     get_axis_list,
     get_latitude_bounds_key,
     get_latitude_key,

pcmdi_metrics/io/xcdat_xarray_dataset_io.py → pcmdi_metrics/io/xcdat_dataset_io.py

@@ -3,6 +3,26 @@
 import xarray as xr
 import xcdat as xc
 
+# Internal function
+
+
+def _find_key(
+    ds: Union[xr.Dataset, xr.DataArray], axis: str, potential_names: list
+) -> str:
+    try:
+        key = xc.get_dim_keys(ds, axis)
+    except Exception:
+        axes = get_axis_list(ds)
+        key_candidates = [k for k in axes if k.lower() in potential_names]
+        if len(key_candidates) > 0:
+            key = key_candidates[0]
+        else:
+            key_candidates = [k for k in axes if k.lower() in potential_names]
+            if len(key_candidates) > 0:
+                key = key_candidates[0]
+    return key
+
+
 # Retrieve coordinate key names
 
 
@@ -11,31 +31,26 @@ def get_axis_list(ds: Union[xr.Dataset, xr.DataArray]) -> list[str]:
     return axes
 
 
+def get_data_list(ds: Union[xr.Dataset, xr.DataArray]) -> list[str]:
+    return list(ds.data_vars.keys())
+
+
 def get_time_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
-    try:
-        time_key = xc.get_dim_keys(ds, "T")
-    except Exception:
-        axes = get_axis_list(ds)
-        time_key = [k for k in axes if k.lower() in ["time"]][0]
-    return time_key
+    axis = "T"
+    potential_names = ["time", "t"]
+    return _find_key(ds, axis, potential_names)
 
 
 def get_latitude_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
-    try:
-        lat_key = xc.get_dim_keys(ds, "Y")
-    except Exception:
-        axes = get_axis_list(ds)
-        lat_key = [k for k in axes if k.lower() in ["lat", "latitude"]][0]
-    return lat_key
+    axis = "Y"
+    potential_names = ["lat", "latitude"]
+    return _find_key(ds, axis, potential_names)
 
 
 def get_longitude_key(ds: Union[xr.Dataset, xr.DataArray]) -> str:
-    try:
-        lon_key = xc.get_dim_keys(ds, "X")
-    except Exception:
-        axes = get_axis_list(ds)
-        lon_key = [k for k in axes if k.lower() in ["lon", "longitude"]][0]
-    return lon_key
+    axis = "X"
+    potential_names = ["lon", "longitude"]
+    return _find_key(ds, axis, potential_names)
 
 
 # Retrieve bounds key names