Duckarray tests for constructors and properties

xarray/testing/duckarrays.py

@@ -0,0 +1,145 @@
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+import hypothesis.extra.numpy as npst
+import hypothesis.strategies as st
+import numpy as np
+import numpy.testing as npt
+import pytest
+from hypothesis import given, note
+
+import xarray as xr
+import xarray.testing.strategies as xrst
+from xarray.core.types import T_DuckArray
+from xarray.testing.assertions import assert_identical
+
+if TYPE_CHECKING:
+    from xarray.core.types import _DTypeLikeNested, _ShapeLike
+
+
+__all__ = [
+    "ConstructorTests",
+    "ReduceTests",
+]
+
+
+class ArrayConstructorChecksMixin:
+    """Mixin for checking results of Variable/DataArray constructors."""
+
+    def check(self, var, arr):
+        self.check_types(var, arr)
+        self.check_values(var, arr)
+        self.check_attributes(var, arr)
+
+    def check_types(self, var, arr):
+        # test type of wrapped array
+        assert isinstance(
+            var.data, type(arr)
+        ), f"found {type(var.data)}, expected {type(arr)}"
+
+    def check_attributes(self, var, arr):
+        # test ndarray attributes are exposed correctly
+        assert var.ndim == arr.ndim
+        assert var.shape == arr.shape
+        assert var.dtype == arr.dtype
+        assert var.size == arr.size
+
+    def check_values(self, var, arr):
+        # test coercion to numpy
+        npt.assert_equal(var.to_numpy(), np.asarray(arr))
+
+
+class ConstructorTests(ArrayConstructorChecksMixin):
+    shapes = npst.array_shapes()
+    dtypes = xrst.supported_dtypes()
+
+    @staticmethod
+    @abstractmethod
+    def array_strategy_fn(
+        *, shape: "_ShapeLike", dtype: "_DTypeLikeNested"
+    ) -> st.SearchStrategy[T_DuckArray]:
+        # TODO can we just make this an attribute?
+        ...
+
+    @given(st.data())
+    def test_construct_variable(self, data) -> None:
+        shape = data.draw(self.shapes)
+        dtype = data.draw(self.dtypes)
+        arr = data.draw(self.array_strategy_fn(shape=shape, dtype=dtype))
+
+        dim_names = data.draw(
+            xrst.dimension_names(min_dims=len(shape), max_dims=len(shape))
+        )
+        var = xr.Variable(data=arr, dims=dim_names)
+
+        self.check(var, arr)
+
+
+def is_real_floating(dtype):
+    return np.issubdtype(dtype, np.number) and np.issubdtype(dtype, np.floating)
+
+
+class ReduceTests:
+    dtypes = xrst.supported_dtypes()
+
+    @staticmethod
+    @abstractmethod
+    def array_strategy_fn(
+        *, shape: "_ShapeLike", dtype: "_DTypeLikeNested"
+    ) -> st.SearchStrategy[T_DuckArray]:
+        # TODO can we just make this an attribute?
+        ...
+
+    def check_reduce(self, var, op, dim, *args, **kwargs):
+        actual = getattr(var, op)(dim=dim, *args, **kwargs)
+
+        data = np.asarray(var.data)
+        expected = getattr(var.copy(data=data), op)(*args, **kwargs)
+
+        # create expected result (using nanmean because arrays with Nans will be generated)
+        reduce_axes = tuple(var.get_axis_num(d) for d in dim)
+        data = np.asarray(var.data)
+        expected = getattr(var.copy(data=data), op)(*args, axis=reduce_axes, **kwargs)
+
+        note(f"actual:\n{actual}")
+        note(f"expected:\n{expected}")
+
+        assert_identical(actual, expected)
+
+    @pytest.mark.parametrize(
+        "method, dtype_assumption",
+        (
+            ("all", lambda x: True),  # should work for any dtype
+            ("any", lambda x: True),  # should work for any dtype
+            # "cumprod",  # not in array API
+            # "cumsum",  # not in array API
+            ("max", is_real_floating),  # only in array API for real numeric dtypes
+            # "median",  # not in array API
+            ("min", is_real_floating),  # only in array API for real numeric dtypes
+            ("prod", is_real_floating),  # only in array API for real numeric dtypes
+            # "std",  # TypeError: std() got an unexpected keyword argument 'ddof'
+            ("sum", is_real_floating),  # only in array API for real numeric dtypes
+            # "var",  # TypeError: std() got an unexpected keyword argument 'ddof'
+        ),
+    )
+    @given(st.data())
+    def test_reduce_variable(self, method, dtype_assumption, data):
+        """
+        Test that the reduction applied to an xarray Variable is always equal
+        to the same reduction applied to the underlying array.
+        """
+
+        narrowed_dtypes = self.dtypes.filter(dtype_assumption)
+
+        var = data.draw(
+            xrst.variables(
+                array_strategy_fn=self.array_strategy_fn,
+                dims=xrst.dimension_names(min_dims=1),
+                dtype=narrowed_dtypes,
+            )
+        )
+
+        # specify arbitrary reduction along at least one dimension
+        reduce_dims = data.draw(xrst.unique_subset_of(var.dims, min_size=1))
+
+        self.check_reduce(var, method, dim=reduce_dims)

xarray/testing/utils.py

@@ -0,0 +1,10 @@
+import warnings
+from contextlib import contextmanager
+
+
+@contextmanager
+def suppress_warning(category, message=""):
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=category, message=message)
+
+        yield

xarray/tests/duckarrays/test_array_api.py

@@ -0,0 +1,45 @@
+from typing import TYPE_CHECKING
+
+import hypothesis.strategies as st
+from hypothesis.extra.array_api import make_strategies_namespace
+
+from xarray.core.types import T_DuckArray
+from xarray.testing import duckarrays
+from xarray.testing.utils import suppress_warning
+from xarray.tests import _importorskip
+
+if TYPE_CHECKING:
+    from xarray.core.types import _DTypeLikeNested, _ShapeLike
+
+
+# ignore the warning that the array_api is experimental raised by numpy
+with suppress_warning(
+    UserWarning, "The numpy.array_api submodule is still experimental. See NEP 47."
+):
+    _importorskip("numpy", "1.26.0")
+    import numpy.array_api as nxp
+
+
+nxps = make_strategies_namespace(nxp)
+
+
+class TestConstructors(duckarrays.ConstructorTests):
+    dtypes = nxps.scalar_dtypes()
+
+    @staticmethod
+    def array_strategy_fn(
+        shape: "_ShapeLike",
+        dtype: "_DTypeLikeNested",
+    ) -> st.SearchStrategy[T_DuckArray]:
+        return nxps.arrays(shape=shape, dtype=dtype)
+
+
+class TestReductions(duckarrays.ReduceTests):
+    dtypes = nxps.scalar_dtypes()
+
+    @staticmethod
+    def array_strategy_fn(
+        shape: "_ShapeLike",
+        dtype: "_DTypeLikeNested",
+    ) -> st.SearchStrategy[T_DuckArray]:
+        return nxps.arrays(shape=shape, dtype=dtype)

xarray/tests/duckarrays/test_numpy.py

@@ -0,0 +1,33 @@
+from typing import TYPE_CHECKING
+
+import hypothesis.extra.numpy as npst
+import hypothesis.strategies as st
+
+from xarray.core.types import T_DuckArray
+from xarray.testing import duckarrays
+from xarray.testing.strategies import supported_dtypes
+
+if TYPE_CHECKING:
+    from xarray.core.types import _DTypeLikeNested, _ShapeLike
+
+
+class TestConstructors(duckarrays.ConstructorTests):
+    dtypes = supported_dtypes()
+
+    @staticmethod
+    def array_strategy_fn(
+        shape: "_ShapeLike",
+        dtype: "_DTypeLikeNested",
+    ) -> st.SearchStrategy[T_DuckArray]:
+        return npst.arrays(shape=shape, dtype=dtype)
+
+
+class TestReductions(duckarrays.ReduceTests):
+    dtypes = supported_dtypes()
+
+    @staticmethod
+    def array_strategy_fn(
+        shape: "_ShapeLike",
+        dtype: "_DTypeLikeNested",
+    ) -> st.SearchStrategy[T_DuckArray]:
+        return npst.arrays(shape=shape, dtype=dtype)

xarray/tests/duckarrays/test_sparse.py

@@ -0,0 +1,76 @@
+from typing import TYPE_CHECKING
+
+import hypothesis.extra.numpy as npst
+import hypothesis.strategies as st
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+import xarray.testing.strategies as xrst
+from xarray.testing import duckarrays
+from xarray.tests import _importorskip
+
+if TYPE_CHECKING:
+    from xarray.core.types import _DTypeLikeNested, _ShapeLike
+
+
+_importorskip("sparse")
+import sparse
+
+
+@pytest.fixture(autouse=True)
+def disable_bottleneck():
+    from xarray import set_options
+
+    with set_options(use_bottleneck=False):
+        yield
+
+
+# sparse does not support float16
+sparse_dtypes = xrst.supported_dtypes().filter(
+    lambda dtype: (not np.issubdtype(dtype, np.float16))
+)
+
+
+@st.composite
+def sparse_arrays_fn(
+    draw: st.DrawFn,
+    *,
+    shape: "_ShapeLike",
+    dtype: "_DTypeLikeNested",
+) -> sparse.COO:
+    """When called generates an arbitrary sparse.COO array of the given shape and dtype."""
+    np_arr = draw(npst.arrays(dtype, shape))
+
+    def to_sparse(arr: np.ndarray) -> sparse.COO:
+        if arr.ndim == 0:
+            return arr
+
+        return sparse.COO.from_numpy(arr)
+
+    return to_sparse(np_arr)
+
+
+class TestConstructors(duckarrays.ConstructorTests):
+    dtypes = sparse_dtypes()
+
+    @staticmethod
+    def array_strategy_fn(
+        shape: "_ShapeLike",
+        dtype: "_DTypeLikeNested",
+    ) -> st.SearchStrategy[sparse.COO]:
+        return sparse_arrays_fn
+
+    def check_values(self, var, arr):
+        npt.assert_equal(var.to_numpy(), arr.todense())
+
+
+class TestReductions(duckarrays.ReduceTests):
+    dtypes = nxps.scalar_dtypes()
+
+    @staticmethod
+    def array_strategy_fn(
+        shape: "_ShapeLike",
+        dtype: "_DTypeLikeNested",
+    ) -> st.SearchStrategy[T_DuckArray]:
+        return nxps.arrays(shape=shape, dtype=dtype)

xarray/tests/test_units.py

@@ -92,11 +92,14 @@ def array_strip_units(array):
 
 
 def array_attach_units(data, unit):
+    if unit is None or (isinstance(unit, int) and unit == 1):
+        return data
+
     if isinstance(data, Quantity):
         raise ValueError(f"cannot attach unit {unit} to quantity {data}")
 
     try:
-        quantity = data * unit
+        quantity = unit._REGISTRY.Quantity(data, unit)
     except np.core._exceptions.UFuncTypeError:
         if isinstance(unit, unit_registry.Unit):
             raise
@@ -181,36 +184,40 @@ def attach_units(obj, units):
         return array_attach_units(obj, units)
 
     if isinstance(obj, xr.Dataset):
-        data_vars = {
-            name: attach_units(value, units) for name, value in obj.data_vars.items()
+        variables = {
+            name: attach_units(value, {None: units.get(name)})
+            for name, value in obj.variables.items()
         }
-
         coords = {
-            name: attach_units(value, units) for name, value in obj.coords.items()
+            name: var for name, var in variables.items() if name in obj._coord_names
+        }
+        data_vars = {
+            name: var for name, var in variables.items() if name not in obj._coord_names
         }
-
         new_obj = xr.Dataset(data_vars=data_vars, coords=coords, attrs=obj.attrs)
     elif isinstance(obj, xr.DataArray):
         # try the array name, "data" and None, then fall back to dimensionless
-        data_units = units.get(obj.name, None) or units.get(None, None) or 1
-
-        data = array_attach_units(obj.data, data_units)
+        units = units.copy()
+        THIS_ARRAY = xr.core.dataarray._THIS_ARRAY
+        unset = object()
+        if obj.name in units:
+            name = obj.name
+        elif None in units:
+            name = None
+        else:
+            name = unset
 
-        coords = {
-            name: (
-                (value.dims, array_attach_units(value.data, units.get(name) or 1))
-                if name in units
-                else (value.dims, value.data)
-            )
-            for name, value in obj.coords.items()
-        }
-        dims = obj.dims
-        attrs = obj.attrs
+        if name is not unset:
+            units[THIS_ARRAY] = units.pop(name)
 
-        new_obj = xr.DataArray(
-            name=obj.name, data=data, coords=coords, attrs=attrs, dims=dims
-        )
+        ds = obj._to_temp_dataset()
+        attached = attach_units(ds, units)
+        new_obj = obj._from_temp_dataset(attached, name=obj.name)
     else:
+        if isinstance(obj, xr.IndexVariable):
+            # no units for index variables
+            return obj
+
         data_units = units.get("data", None) or units.get(None, None) or 1
 
         data = array_attach_units(obj.data, data_units)