types, missing parameters

Author: amalia-k510

src/fast_array_utils/stats/_is_constant.py

@@ -19,7 +19,7 @@
 
 @singledispatch
 def is_constant_(
-    a: NDArray[Any] | types.CSBase | types.CupyArray | types.DaskArray,
+    a: NDArray[Any] | types.CSBase | types.CupyArray | types.DaskArray | types.HasArrayNamespace,
     /,
     *,
     axis: Literal[0, 1] | None = None,

src/fast_array_utils/stats/_mean.py

@@ -18,7 +18,7 @@
 
 
 def mean_(
-    x: CpuArray | GpuArray | DiskArray | types.DaskArray,
+    x: CpuArray | GpuArray | DiskArray | types.DaskArray | types.HasArrayNamespace,
     /,
     *,
     axis: Literal[0, 1] | None = None,

src/fast_array_utils/stats/_mean_var.py

@@ -21,7 +21,7 @@
 
 @no_type_check  # mypy is extremely confused
 def mean_var_(
-    x: CpuArray | GpuArray | types.DaskArray,
+    x: CpuArray | GpuArray | types.DaskArray | types.HasArrayNamespace,
     /,
     *,
     axis: Literal[0, 1] | None = None,

src/fast_array_utils/stats/_typing.py

@@ -28,22 +28,27 @@ class StatFunNoDtype(Protocol):
     __name__: str
 
     def __call__(
-        self, x: CpuArray | GpuArray | DiskArray | types.DaskArray, /, *, axis: Literal[0, 1] | None = None, keep_cupy_as_array: bool = False
-    ) -> types.DaskArray: ...
+        self,
+        x: CpuArray | GpuArray | DiskArray | types.DaskArray | types.HasArrayNamespace,
+        /,
+        *,
+        axis: Literal[0, 1] | None = None,
+        keep_cupy_as_array: bool = False,
+    ) -> types.DaskArray | types.HasArrayNamespace: ...
 
 
 class StatFunDtype(Protocol):
     __name__: str
 
     def __call__(
         self,
-        x: CpuArray | GpuArray | DiskArray | types.DaskArray,
+        x: CpuArray | GpuArray | DiskArray | types.DaskArray | types.HasArrayNamespace,
         /,
         *,
         axis: Literal[0, 1] | None = None,
         dtype: DTypeLike | None = None,
         keep_cupy_as_array: bool = False,
-    ) -> NDArray[Any] | types.CupyArray | np.number[Any] | types.DaskArray: ...
+    ) -> NDArray[Any] | types.CupyArray | np.number[Any] | types.DaskArray | types.HasArrayNamespace: ...
 
 
 NoDtypeOps = Literal["max", "min"]

src/fast_array_utils/types.py

@@ -124,4 +124,7 @@
 class HasArrayNamespace(Protocol):
     """An array object compatible with the Python array API standard."""
 
+    ndim: int
+    shape: tuple[int, ...]
+
     def __array_namespace__(self, /, *, api_version: str | None = None) -> ModuleType: ...