Tests for variables with units

xarray/tests/test_units.py

@@ -7,12 +7,13 @@
 import xarray as xr
 from xarray.core import formatting
 from xarray.core.npcompat import IS_NEP18_ACTIVE
+from .test_variable import VariableSubclassobjects
 
 pint = pytest.importorskip("pint")
 DimensionalityError = pint.errors.DimensionalityError
 
 
-unit_registry = pint.UnitRegistry()
+unit_registry = pint.UnitRegistry(force_ndarray=True)
 Quantity = unit_registry.Quantity
 
 pytestmark = [
@@ -113,6 +114,10 @@ def extract_units(obj):
         }
 
         units = {**vars_units, **coords_units}
+    elif isinstance(obj, xr.Variable):
+        vars_units = {None: array_extract_units(obj.data)}
+
+        units = {**vars_units}
     elif isinstance(obj, Quantity):
         vars_units = {None: array_extract_units(obj)}
 
@@ -148,6 +153,9 @@ def strip_units(obj):
         new_obj = xr.DataArray(
             name=strip_units(obj.name), data=data, coords=coords, dims=obj.dims
         )
+    elif isinstance(obj, xr.Variable):
+        data = array_strip_units(obj.data)
+        new_obj = obj.copy(data=data)
     elif isinstance(obj, unit_registry.Quantity):
         new_obj = obj.magnitude
     elif isinstance(obj, (list, tuple)):
@@ -159,7 +167,7 @@ def strip_units(obj):
 
 
 def attach_units(obj, units):
-    if not isinstance(obj, (xr.DataArray, xr.Dataset)):
+    if not isinstance(obj, (xr.DataArray, xr.Dataset, xr.Variable)):
         return array_attach_units(obj, units.get("data", 1))
 
     if isinstance(obj, xr.Dataset):
@@ -172,7 +180,7 @@ def attach_units(obj, units):
         }
 
         new_obj = xr.Dataset(data_vars=data_vars, coords=coords, attrs=obj.attrs)
-    else:
+    elif isinstance(obj, xr.DataArray):
         # try the array name, "data" and None, then fall back to dimensionless
         data_units = (
             units.get(obj.name, None)
@@ -198,6 +206,16 @@ def attach_units(obj, units):
         new_obj = xr.DataArray(
             name=obj.name, data=data, coords=coords, attrs=attrs, dims=dims
         )
+    else:
+        data_units = (
+            units.get(obj.name, None)
+            or units.get("data", None)
+            or units.get(None, None)
+            or 1
+        )
+
+        data = array_attach_units(obj.data, data_units)
+        new_obj = obj.copy(data=data)
 
     return new_obj
 
@@ -1245,6 +1263,146 @@ def test_dot_dataarray(dtype):
     assert_equal_with_units(expected, actual)
 
 
+def delete_attrs(*to_delete):
+    def wrapper(cls):
+        for item in to_delete:
+            setattr(cls, item, None)
+
+        return cls
+
+    return wrapper
+
+
+@delete_attrs(
+    "test_getitem_with_mask",
+    "test_getitem_with_mask_nd_indexer",
+    "test_index_0d_string",
+    "test_index_0d_datetime",
+    "test_index_0d_timedelta64",
+    "test_0d_time_data",
+    "test_datetime64_conversion",
+    "test_timedelta64_conversion",
+    "test_pandas_period_index",
+    "test_1d_math",
+    "test_1d_reduce",
+    "test_array_interface",
+    "test___array__",
+    "test_copy_index",
+    "test_concat",
+    "test_concat_number_strings",
+    "test_concat_fixed_len_str",
+    "test_concat_mixed_dtypes",
+    "test_pandas_datetime64_with_tz",
+    "test_multiindex",
+)
+class TestVariable(VariableSubclassobjects):
+    @staticmethod
+    def cls(dims, data, *args, **kwargs):
+        return xr.Variable(
+            dims, unit_registry.Quantity(data, unit_registry.m), *args, **kwargs
+        )
+
+    @pytest.mark.parametrize(
+        "func",
+        (
+            pytest.param(
+                method("all"), marks=pytest.mark.xfail(reason="not implemented by pint")
+            ),
+            pytest.param(
+                method("any"), marks=pytest.mark.xfail(reason="not implemented by pint")
+            ),
+            method("argmax"),
+            method("argmin"),
+            method("argsort"),
+            method("cumprod"),
+            method("cumsum"),
+            method("max"),
+            method("mean"),
+            method("median"),
+            method("min"),
+            pytest.param(
+                method("prod"),
+                marks=pytest.mark.xfail(reason="not implemented by pint"),
+            ),
+            method("std"),
+            method("sum"),
+            method("var"),
+        ),
+        ids=repr,
+    )
+    def test_aggregation(self, func, dtype):
+        array = np.linspace(0, 1, 10).astype(dtype) * (
+            unit_registry.m if func.name != "cumprod" else unit_registry.dimensionless
+        )
+        variable = xr.Variable("x", array)
+
+        units = extract_units(func(array))
+        expected = attach_units(func(strip_units(variable)), units)
+        actual = func(variable)
+
+        xr.testing.assert_identical(expected, actual)
+
+    @pytest.mark.parametrize(
+        "func",
+        (
+            method("astype", np.float32),
+            method("conj"),
+            method("conjugate"),
+            method("searchsorted", 5),
+            method("clip", min=2, max=7),
+        ),
+        ids=repr,
+    )
+    @pytest.mark.parametrize(
+        "unit,error",
+        (
+            pytest.param(1, DimensionalityError, id="no_unit"),
+            pytest.param(
+                unit_registry.dimensionless, DimensionalityError, id="dimensionless"
+            ),
+            pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
+            pytest.param(unit_registry.cm, None, id="compatible_unit"),
+            pytest.param(unit_registry.m, None, id="identical_unit"),
+        ),
+    )
+    def test_numpy_methods(self, func, unit, error, dtype):
+        array = np.linspace(0, 1, 10).astype(dtype) * unit_registry.m
+        variable = xr.Variable("x", array)
+
+        args = [
+            item * unit if isinstance(item, (int, float, list)) else item
+            for item in func.args
+        ]
+        kwargs = {
+            key: value * unit if isinstance(value, (int, float, list)) else value
+            for key, value in func.kwargs.items()
+        }
+
+        if error is not None and func.name in ("searchsorted", "clip"):
+            with pytest.raises(error):
+                func(variable, *args, **kwargs)
+
+            return
+
+        converted_args = [
+            strip_units(convert_units(item, {None: unit_registry.m})) for item in args
+        ]
+        converted_kwargs = {
+            key: strip_units(convert_units(value, {None: unit_registry.m}))
+            for key, value in kwargs.items()
+        }
+
+        print("running on:", func, args, kwargs)
+        units = extract_units(func(array, *args, **kwargs))
+        expected = attach_units(
+            func(strip_units(variable), *converted_args, **converted_kwargs), units
+        )
+        actual = func(variable, *args, **kwargs)
+
+        assert extract_units(expected) == extract_units(actual)
+        xr.testing.assert_allclose(expected, actual)
+
+
 class TestDataArray:
     @pytest.mark.filterwarnings("error:::pint[.*]")
     @pytest.mark.parametrize(