Skip to content

Commit

Permalink
Add duckarray test for np.array_api (#8391)
Browse files Browse the repository at this point in the history
  • Loading branch information
Illviljan committed Oct 31, 2023
1 parent f63ede9 commit cfe4d71
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 17 deletions.
21 changes: 13 additions & 8 deletions xarray/namedarray/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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__(
Expand Down
34 changes: 25 additions & 9 deletions xarray/tests/test_namedarray.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand All @@ -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}
Expand All @@ -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)

Expand All @@ -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


Expand Down Expand Up @@ -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 == {}
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit cfe4d71

Please sign in to comment.