Merge branch 'add_array_api_mean' of https://github.com/Illviljan/xarray

into add_array_api_mean
pydata · Aug 22, 2024 · 10a8662 · 10a8662
2 parents 82d3af9 + 10a19f4
commit 10a8662
Show file tree

Hide file tree

Showing 8 changed files with 30 additions and 34 deletions.
diff --git a/xarray/namedarray/_array_api/_indexing_functions.py b/xarray/namedarray/_array_api/_indexing_functions.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from xarray.namedarray._array_api._utils import _get_data_namespace, _dims_to_axis
+from xarray.namedarray._array_api._utils import _dims_to_axis, _get_data_namespace
 from xarray.namedarray._typing import (
     Default,
     _default,

diff --git a/xarray/namedarray/_array_api/_info.py b/xarray/namedarray/_array_api/_info.py
@@ -2,7 +2,6 @@
 
 from typing import TYPE_CHECKING
 
-
 from xarray.namedarray._array_api._dtypes import (
     bool,
     complex64,
@@ -23,10 +22,10 @@
     from types import ModuleType
 
     from xarray.namedarray._typing import (
-        _Device,
         _Capabilities,
-        _DefaultDataTypes,
         _DataTypes,
+        _DefaultDataTypes,
+        _Device,
     )
 
 

diff --git a/xarray/namedarray/_array_api/_searching_functions.py b/xarray/namedarray/_array_api/_searching_functions.py
@@ -3,10 +3,10 @@
 from typing import TYPE_CHECKING
 
 from xarray.namedarray._array_api._utils import (
-    _get_data_namespace,
-    _infer_dims,
     _dims_to_axis,
+    _get_data_namespace,
     _get_remaining_dims,
+    _infer_dims,
 )
 from xarray.namedarray._typing import (
     Default,

diff --git a/xarray/namedarray/_array_api/_sorting_functions.py b/xarray/namedarray/_array_api/_sorting_functions.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from xarray.namedarray._array_api._utils import _get_data_namespace, _dims_to_axis
+from xarray.namedarray._array_api._utils import _dims_to_axis, _get_data_namespace
 from xarray.namedarray._typing import (
     Default,
     _default,

diff --git a/xarray/namedarray/_array_api/_statistical_functions.py b/xarray/namedarray/_array_api/_statistical_functions.py
@@ -3,8 +3,8 @@
 from typing import Any
 
 from xarray.namedarray._array_api._utils import (
-    _get_data_namespace,
     _dims_to_axis,
+    _get_data_namespace,
     _get_remaining_dims,
 )
 from xarray.namedarray._typing import (

diff --git a/xarray/namedarray/_array_api/_utility_functions.py b/xarray/namedarray/_array_api/_utility_functions.py
@@ -3,8 +3,8 @@
 from typing import Any
 
 from xarray.namedarray._array_api._utils import (
-    _get_data_namespace,
     _dims_to_axis,
+    _get_data_namespace,
     _get_remaining_dims,
 )
 from xarray.namedarray._typing import (

diff --git a/xarray/namedarray/_array_api/_utils.py b/xarray/namedarray/_array_api/_utils.py
@@ -1,19 +1,20 @@
 from __future__ import annotations
 
+from collections.abc import Iterable
 from types import ModuleType
-from typing import TYPE_CHECKING, Any, Iterable
+from typing import TYPE_CHECKING, Any
 
 from xarray.namedarray._typing import (
     Default,
     _arrayapi,
-    _default,
-    _DimsLike,
-    _dtype,
     _AxisLike,
-    _Shape,
+    _default,
     _Dim,
     _Dims,
+    _DimsLike,
     _DType,
+    _dtype,
+    _Shape,
     duckarray,
 )
 

diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py
@@ -11,11 +11,11 @@
     Literal,
     Protocol,
     SupportsIndex,
+    TypedDict,
     TypeVar,
     Union,
     overload,
     runtime_checkable,
-    TypedDict,
 )
 
 import numpy as np
@@ -144,25 +144,21 @@ class _FInfo(Protocol):
     },
 )
 
-_DataTypes = TypedDict(
-    "DataTypes",
-    {
-        "bool": _dtype,
-        "float32": _dtype,
-        "float64": _dtype,
-        "complex64": _dtype,
-        "complex128": _dtype,
-        "int8": _dtype,
-        "int16": _dtype,
-        "int32": _dtype,
-        "int64": _dtype,
-        "uint8": _dtype,
-        "uint16": _dtype,
-        "uint32": _dtype,
-        "uint64": _dtype,
-    },
-    total=False,
-)
+
+class _DataTypes(TypedDict, total=False):
+    bool: _dtype
+    float32: _dtype
+    float64: _dtype
+    complex64: _dtype
+    complex128: _dtype
+    int8: _dtype
+    int16: _dtype
+    int32: _dtype
+    int64: _dtype
+    uint8: _dtype
+    uint16: _dtype
+    uint32: _dtype
+    uint64: _dtype
 
 
 @runtime_checkable