diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 10314fe9440..8cfc6931431 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -93,14 +93,6 @@ def shape(self) -> _Shape: def dtype(self) -> _DType_co: ... - @overload - def __array__(self, dtype: None = ..., /) -> np.ndarray[Any, _DType_co]: - ... - - @overload - def __array__(self, dtype: _DType, /) -> np.ndarray[Any, _DType]: - ... - @runtime_checkable class _arrayfunction( @@ -112,6 +104,19 @@ class _arrayfunction( Corresponds to np.ndarray. """ + @overload + def __array__(self, dtype: None = ..., /) -> np.ndarray[Any, _DType_co]: + ... + + @overload + def __array__(self, dtype: _DType, /) -> np.ndarray[Any, _DType]: + ... + + def __array__( + self, dtype: _DType | None = ..., / + ) -> np.ndarray[Any, _DType] | np.ndarray[Any, _DType_co]: + ... + # TODO: Should return the same subclass but with a new dtype generic. # https://github.com/python/typing/issues/548 def __array_ufunc__( diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 6e39a3aa94f..448e8cf819a 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -1,6 +1,7 @@ from __future__ import annotations import copy +import warnings from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Generic, cast, overload @@ -66,13 +67,13 @@ def test_namedarray_init() -> None: expected = np.array([1, 2], dtype=dtype) actual: NamedArray[Any, np.dtype[np.int8]] actual = NamedArray(("x",), expected) - assert np.array_equal(actual.data, expected) + assert np.array_equal(np.asarray(actual.data), expected) with pytest.raises(AttributeError): expected2 = [1, 2] actual2: NamedArray[Any, Any] actual2 = NamedArray(("x",), expected2) # type: ignore[arg-type] - assert np.array_equal(actual2.data, expected2) + assert np.array_equal(np.asarray(actual2.data), expected2) @pytest.mark.parametrize( @@ -101,7 +102,7 @@ def test_from_array( else: actual = from_array(dims, data) - assert np.array_equal(actual.data, expected) + assert np.array_equal(np.asarray(actual.data), expected) def test_from_array_with_masked_array() -> None: @@ -114,7 +115,8 @@ def test_from_array_with_masked_array() -> None: def test_from_array_with_0d_object() -> None: data = np.empty((), dtype=object) data[()] = (10, 12, 12) - np.array_equal(from_array((), data).data, data) + narr = from_array((), data) + np.array_equal(np.asarray(narr.data), data) # TODO: Make xr.core.indexing.ExplicitlyIndexed pass as a subclass of_arrayfunction_or_api @@ -140,7 +142,7 @@ def test_properties() -> None: named_array: NamedArray[Any, Any] named_array = NamedArray(["x", "y"], data, {"key": "value"}) assert named_array.dims == ("x", "y") - assert np.array_equal(named_array.data, data) + assert np.array_equal(np.asarray(named_array.data), data) assert named_array.attrs == {"key": "value"} assert named_array.ndim == 2 assert named_array.sizes == {"x": 2, "y": 5} @@ -162,7 +164,7 @@ def test_attrs() -> None: def test_data(random_inputs: np.ndarray[Any, Any]) -> None: named_array: NamedArray[Any, Any] named_array = NamedArray(["x", "y", "z"], random_inputs) - assert np.array_equal(named_array.data, random_inputs) + assert np.array_equal(np.asarray(named_array.data), random_inputs) with pytest.raises(ValueError): named_array.data = np.random.random((3, 4)).astype(np.float64) @@ -181,11 +183,11 @@ def test_real_and_imag() -> None: named_array = NamedArray(["x"], arr) actual_real: duckarray[Any, np.dtype[np.float64]] = named_array.real.data - assert np.array_equal(actual_real, expected_real) + assert np.array_equal(np.asarray(actual_real), expected_real) assert actual_real.dtype == expected_real.dtype actual_imag: duckarray[Any, np.dtype[np.float64]] = named_array.imag.data - assert np.array_equal(actual_imag, expected_imag) + assert np.array_equal(np.asarray(actual_imag), expected_imag) assert actual_imag.dtype == expected_imag.dtype @@ -214,7 +216,7 @@ def test_0d_object() -> None: named_array = from_array([], (10, 12, 12)) expected_data = np.empty((), dtype=object) expected_data[()] = (10, 12, 12) - assert np.array_equal(named_array.data, expected_data) + assert np.array_equal(np.asarray(named_array.data), expected_data) assert named_array.dims == () assert named_array.sizes == {} @@ -294,6 +296,20 @@ def test_duck_array_typevar(a: duckarray[Any, _DType]) -> duckarray[Any, _DType] test_duck_array_typevar(numpy_a) test_duck_array_typevar(custom_a) + # Test numpy's array api: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + r"The numpy.array_api submodule is still experimental", + category=UserWarning, + ) + import numpy.array_api as nxp + + # TODO: nxp doesn't use dtype typevars, so can only use Any for the moment: + arrayapi_a: duckarray[Any, Any] # duckarray[Any, np.dtype[np.int64]] + arrayapi_a = nxp.asarray([2.1, 4], dtype=np.dtype(np.int64)) + test_duck_array_typevar(arrayapi_a) + def test_new_namedarray() -> None: dtype_float = np.dtype(np.float32)