add .imag and .real properties to NamedArray (#8365)

Author: andersy005

xarray/core/variable.py

@@ -2365,28 +2365,6 @@ def notnull(self, keep_attrs: bool | None = None):
             keep_attrs=keep_attrs,
         )
 
-    @property
-    def real(self):
-        """
-        The real part of the variable.
-
-        See Also
-        --------
-        numpy.ndarray.real
-        """
-        return self._replace(data=self.data.real)
-
-    @property
-    def imag(self):
-        """
-        The imaginary part of the variable.
-
-        See Also
-        --------
-        numpy.ndarray.imag
-        """
-        return self._replace(data=self.data.imag)
-
     def __array_wrap__(self, obj, context=None):
         return Variable(self.dims, obj)
 

xarray/namedarray/core.py

@@ -24,16 +24,11 @@
 from xarray.namedarray._typing import (
     _arrayfunction_or_api,
     _chunkedarray,
-    _DType,
     _DType_co,
     _ScalarType_co,
     _ShapeType_co,
 )
-from xarray.namedarray.utils import (
-    _default,
-    is_duck_dask_array,
-    to_0d_object_array,
-)
+from xarray.namedarray.utils import _default, is_duck_dask_array, to_0d_object_array
 
 if TYPE_CHECKING:
     from numpy.typing import ArrayLike, NDArray
@@ -46,6 +41,7 @@
         _Dim,
         _Dims,
         _DimsLike,
+        _DType,
         _IntOrUnknown,
         _ScalarType,
         _Shape,
@@ -516,6 +512,28 @@ def data(self, data: duckarray[Any, _DType_co]) -> None:
         self._check_shape(data)
         self._data = data
 
+    @property
+    def imag(self) -> Self:
+        """
+        The imaginary part of the array.
+
+        See Also
+        --------
+        numpy.ndarray.imag
+        """
+        return self._replace(data=self.data.imag)  # type: ignore
+
+    @property
+    def real(self) -> Self:
+        """
+        The real part of the array.
+
+        See Also
+        --------
+        numpy.ndarray.real
+        """
+        return self._replace(data=self.data.real)  # type: ignore
+
     def __dask_tokenize__(self) -> Hashable:
         # Use v.data, instead of v._data, in order to cope with the wrappers
         # around NetCDF and the like

xarray/tests/test_namedarray.py

@@ -8,11 +8,7 @@
 import pytest
 
 from xarray.core.indexing import ExplicitlyIndexed
-from xarray.namedarray._typing import (
-    _arrayfunction_or_api,
-    _DType_co,
-    _ShapeType_co,
-)
+from xarray.namedarray._typing import _arrayfunction_or_api, _DType_co, _ShapeType_co
 from xarray.namedarray.core import NamedArray, from_array
 from xarray.namedarray.utils import _default
 
@@ -171,6 +167,16 @@ def test_data(random_inputs: np.ndarray[Any, Any]) -> None:
         named_array.data = np.random.random((3, 4)).astype(np.float64)
 
 
+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_imag = -np.arange(3)
+    assert np.array_equal(named_array.imag.data, expected_imag)
+
+
 # Additional tests as per your original class-based code
 @pytest.mark.parametrize(
     "data, dtype",