From dd5eb512286135b9a542516e72bff3a4da32fd84 Mon Sep 17 00:00:00 2001
From: Illviljan <14371165+Illviljan@users.noreply.github.com>
Date: Sun, 22 Oct 2023 02:52:02 +0200
Subject: [PATCH] Use namedarray repr in _array_api docstrings (#8355)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 xarray/core/formatting.py       |  2 +-
 xarray/namedarray/_array_api.py | 46 +++++++++++++++++++++++++--------
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/xarray/core/formatting.py b/xarray/core/formatting.py
index 96a767f95ac..561b8d3cc0d 100644
--- a/xarray/core/formatting.py
+++ b/xarray/core/formatting.py
@@ -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
diff --git a/xarray/namedarray/_array_api.py b/xarray/namedarray/_array_api.py
index 69f97305686..e205c4d4efe 100644
--- a/xarray/namedarray/_array_api.py
+++ b/xarray/namedarray/_array_api.py
@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+import warnings
 from types import ModuleType
 from typing import Any
 
@@ -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(
@@ -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]]:
@@ -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)
@@ -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