Skip to content

Commit dd5eb51

Browse files
Use namedarray repr in _array_api docstrings (#8355)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent b0bb86e commit dd5eb51

File tree

2 files changed

+36
-12
lines changed

2 files changed

+36
-12
lines changed

xarray/core/formatting.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -632,7 +632,7 @@ def short_data_repr(array):
632632
return short_array_repr(array)
633633
elif is_duck_array(internal_data):
634634
return limit_lines(repr(array.data), limit=40)
635-
elif array._in_memory:
635+
elif getattr(array, "_in_memory", None):
636636
return short_array_repr(array)
637637
else:
638638
# internal xarray array type

xarray/namedarray/_array_api.py

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
from __future__ import annotations
2+
3+
import warnings
14
from types import ModuleType
25
from typing import Any
36

@@ -13,12 +16,23 @@
1316
)
1417
from xarray.namedarray.core import NamedArray
1518

19+
with warnings.catch_warnings():
20+
warnings.filterwarnings(
21+
"ignore",
22+
r"The numpy.array_api submodule is still experimental",
23+
category=UserWarning,
24+
)
25+
import numpy.array_api as nxp # noqa: F401
26+
1627

1728
def _get_data_namespace(x: NamedArray[Any, Any]) -> ModuleType:
1829
if isinstance(x._data, _arrayapi):
1930
return x._data.__array_namespace__()
20-
else:
21-
return np
31+
32+
return np
33+
34+
35+
# %% Creation Functions
2236

2337

2438
def astype(
@@ -49,18 +63,25 @@ def astype(
4963
5064
Examples
5165
--------
52-
>>> narr = NamedArray(("x",), np.array([1.5, 2.5]))
53-
>>> astype(narr, np.dtype(int)).data
54-
array([1, 2])
66+
>>> narr = NamedArray(("x",), nxp.asarray([1.5, 2.5]))
67+
>>> narr
68+
<xarray.NamedArray (x: 2)>
69+
Array([1.5, 2.5], dtype=float64)
70+
>>> astype(narr, np.dtype(np.int32))
71+
<xarray.NamedArray (x: 2)>
72+
Array([1, 2], dtype=int32)
5573
"""
5674
if isinstance(x._data, _arrayapi):
5775
xp = x._data.__array_namespace__()
58-
return x._new(data=xp.astype(x, dtype, copy=copy))
76+
return x._new(data=xp.astype(x._data, dtype, copy=copy))
5977

6078
# np.astype doesn't exist yet:
6179
return x._new(data=x._data.astype(dtype, copy=copy)) # type: ignore[attr-defined]
6280

6381

82+
# %% Elementwise Functions
83+
84+
6485
def imag(
6586
x: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], / # type: ignore[type-var]
6687
) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]:
@@ -83,8 +104,9 @@ def imag(
83104
84105
Examples
85106
--------
86-
>>> narr = NamedArray(("x",), np.array([1 + 2j, 2 + 4j]))
87-
>>> imag(narr).data
107+
>>> narr = NamedArray(("x",), np.asarray([1.0 + 2j, 2 + 4j])) # TODO: Use nxp
108+
>>> imag(narr)
109+
<xarray.NamedArray (x: 2)>
88110
array([2., 4.])
89111
"""
90112
xp = _get_data_namespace(x)
@@ -114,9 +136,11 @@ def real(
114136
115137
Examples
116138
--------
117-
>>> narr = NamedArray(("x",), np.array([1 + 2j, 2 + 4j]))
118-
>>> real(narr).data
139+
>>> narr = NamedArray(("x",), np.asarray([1.0 + 2j, 2 + 4j])) # TODO: Use nxp
140+
>>> real(narr)
141+
<xarray.NamedArray (x: 2)>
119142
array([1., 2.])
120143
"""
121144
xp = _get_data_namespace(x)
122-
return x._new(data=xp.real(x._data))
145+
out = x._new(data=xp.real(x._data))
146+
return out

0 commit comments

Comments
 (0)