Skip to content

Commit

Permalink
fix NamedArray.imag and NamedArray.real typing info
Browse files Browse the repository at this point in the history
  • Loading branch information
andersy005 committed Oct 24, 2023
1 parent ccc8f99 commit 3dcd501
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 10 deletions.
14 changes: 10 additions & 4 deletions xarray/namedarray/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
_DType_co,
_ScalarType_co,
_ShapeType_co,
_SupportsImag,
_SupportsReal,
)
from xarray.namedarray.utils import _default, is_duck_dask_array, to_0d_object_array

Expand Down Expand Up @@ -513,26 +515,30 @@ def data(self, data: duckarray[Any, _DType_co]) -> None:
self._data = data

@property
def imag(self) -> Self:
def imag(
self: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], # type: ignore[type-var]
) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]:
"""
The imaginary part of the array.
See Also
--------
numpy.ndarray.imag
"""
return self._replace(data=self.data.imag) # type: ignore
return self._new(data=self.data.imag)

@property
def real(self) -> Self:
def real(
self: NamedArray[_ShapeType, np.dtype[_SupportsReal[_ScalarType]]], # type: ignore[type-var]
) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]:
"""
The real part of the array.
See Also
--------
numpy.ndarray.real
"""
return self._replace(data=self.data.real) # type: ignore
return self._new(data=self.data.real)

def __dask_tokenize__(self) -> Hashable:
# Use v.data, instead of v._data, in order to cope with the wrappers
Expand Down
24 changes: 18 additions & 6 deletions xarray/tests/test_namedarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,13 +168,25 @@ def test_data(random_inputs: np.ndarray[Any, Any]) -> None:


def test_real_and_imag() -> None:
named_array: NamedArray[Any, Any]
named_array = NamedArray(["x"], np.arange(3) - 1j * np.arange(3))
expected_real = np.arange(3)
assert np.array_equal(named_array.real.data, expected_real)
expected_real: np.ndarray[Any, np.dtype[np.float64]]
expected_real = np.arange(3, dtype=np.float64)

expected_imag: np.ndarray[Any, np.dtype[np.float64]]
expected_imag = -np.arange(3, dtype=np.float64)

arr: np.ndarray[Any, np.dtype[np.complex128]]
arr = expected_real + 1j * expected_imag

named_array: NamedArray[Any, np.dtype[np.complex128]]
named_array = NamedArray(["x"], arr)

actual_real: np.ndarray[Any, np.dtype[np.float64]] = named_array.real.data
assert np.array_equal(actual_real, expected_real)
assert actual_real.dtype == expected_real.dtype

expected_imag = -np.arange(3)
assert np.array_equal(named_array.imag.data, expected_imag)
actual_imag: np.ndarray[Any, np.dtype[np.float64]] = named_array.imag.data
assert np.array_equal(actual_imag, expected_imag)
assert actual_imag.dtype == expected_imag.dtype


# Additional tests as per your original class-based code
Expand Down

0 comments on commit 3dcd501

Please sign in to comment.