Skip to content

Commit

Permalink
Use namedarray repr in _array_api docstrings (#8355)
Browse files Browse the repository at this point in the history
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
Illviljan and pre-commit-ci[bot] authored Oct 22, 2023
1 parent b0bb86e commit dd5eb51
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 12 deletions.
2 changes: 1 addition & 1 deletion xarray/core/formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -632,7 +632,7 @@ def short_data_repr(array):
return short_array_repr(array)
elif is_duck_array(internal_data):
return limit_lines(repr(array.data), limit=40)
elif array._in_memory:
elif getattr(array, "_in_memory", None):
return short_array_repr(array)
else:
# internal xarray array type
Expand Down
46 changes: 35 additions & 11 deletions xarray/namedarray/_array_api.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
from __future__ import annotations

import warnings
from types import ModuleType
from typing import Any

Expand All @@ -13,12 +16,23 @@
)
from xarray.namedarray.core import NamedArray

with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
r"The numpy.array_api submodule is still experimental",
category=UserWarning,
)
import numpy.array_api as nxp # noqa: F401


def _get_data_namespace(x: NamedArray[Any, Any]) -> ModuleType:
if isinstance(x._data, _arrayapi):
return x._data.__array_namespace__()
else:
return np

return np


# %% Creation Functions


def astype(
Expand Down Expand Up @@ -49,18 +63,25 @@ def astype(
Examples
--------
>>> narr = NamedArray(("x",), np.array([1.5, 2.5]))
>>> astype(narr, np.dtype(int)).data
array([1, 2])
>>> narr = NamedArray(("x",), nxp.asarray([1.5, 2.5]))
>>> narr
<xarray.NamedArray (x: 2)>
Array([1.5, 2.5], dtype=float64)
>>> astype(narr, np.dtype(np.int32))
<xarray.NamedArray (x: 2)>
Array([1, 2], dtype=int32)
"""
if isinstance(x._data, _arrayapi):
xp = x._data.__array_namespace__()
return x._new(data=xp.astype(x, dtype, copy=copy))
return x._new(data=xp.astype(x._data, dtype, copy=copy))

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


# %% Elementwise Functions


def imag(
x: NamedArray[_ShapeType, np.dtype[_SupportsImag[_ScalarType]]], / # type: ignore[type-var]
) -> NamedArray[_ShapeType, np.dtype[_ScalarType]]:
Expand All @@ -83,8 +104,9 @@ def imag(
Examples
--------
>>> narr = NamedArray(("x",), np.array([1 + 2j, 2 + 4j]))
>>> imag(narr).data
>>> narr = NamedArray(("x",), np.asarray([1.0 + 2j, 2 + 4j])) # TODO: Use nxp
>>> imag(narr)
<xarray.NamedArray (x: 2)>
array([2., 4.])
"""
xp = _get_data_namespace(x)
Expand Down Expand Up @@ -114,9 +136,11 @@ def real(
Examples
--------
>>> narr = NamedArray(("x",), np.array([1 + 2j, 2 + 4j]))
>>> real(narr).data
>>> narr = NamedArray(("x",), np.asarray([1.0 + 2j, 2 + 4j])) # TODO: Use nxp
>>> real(narr)
<xarray.NamedArray (x: 2)>
array([1., 2.])
"""
xp = _get_data_namespace(x)
return x._new(data=xp.real(x._data))
out = x._new(data=xp.real(x._data))
return out

0 comments on commit dd5eb51

Please sign in to comment.