undo disk array

Author: flying-sheep

src/fast_array_utils/stats/_generic_ops.py

@@ -87,7 +87,7 @@ def _generic_op_cs(
     if TYPE_CHECKING:  # scipy-stubs thinks e.g. "int64" is invalid, which isn’t true
         assert isinstance(dtype, np.dtype | type | None)
     # convert to array so dimensions collapse as expected
-    x = (sp.csr_array if x.format == "csr" else sp.csc_array)(x, **_dtype_kw(dtype, op))  # type: ignore[call-overload]
+    x = (sp.csr_array if x.format == "csr" else sp.csc_array)(x, **_dtype_kw(dtype, op))  # type: ignore[arg-type]
     return cast("NDArray[Any] | np.number[Any]", getattr(x, op)(axis=axis))
 
 

src/fast_array_utils/stats/_mean_var.py

@@ -15,12 +15,12 @@
 
     from numpy.typing import NDArray
 
-    from ..typing import CpuArray, DiskArray, GpuArray
+    from ..typing import CpuArray, GpuArray
 
 
 @no_type_check  # mypy is extremely confused
 def mean_var_(
-    x: CpuArray | GpuArray | DiskArray | types.DaskArray,
+    x: CpuArray | GpuArray | types.DaskArray,
     /,
     *,
     axis: Literal[0, 1] | None = None,

src/fast_array_utils/stats/_power.py

@@ -14,10 +14,10 @@
 
     from numpy.typing import DTypeLike
 
-    from fast_array_utils.typing import CpuArray, DiskArray, GpuArray
+    from fast_array_utils.typing import CpuArray, GpuArray
 
     # All supported array types except for disk ones and CSDataset
-    Array: TypeAlias = CpuArray | GpuArray | DiskArray | types.DaskArray
+    Array: TypeAlias = CpuArray | GpuArray | types.DaskArray
 
     _Arr = TypeVar("_Arr", bound=Array)
     _Mat = TypeVar("_Mat", bound=types.CSBase | types.CupyCSMatrix)
@@ -33,7 +33,7 @@ def power(x: _Arr, n: int, /, dtype: DTypeLike | None = None) -> _Arr:
 def _power(x: Array, n: int, /, dtype: DTypeLike | None = None) -> Array:
     if TYPE_CHECKING:
         assert not isinstance(x, types.DaskArray | types.CSBase | types.CupyCSMatrix)
-    return np.power(x, n, dtype=dtype)  # type: ignore[operator]
+    return x**n if dtype is None else np.power(x, n, dtype=dtype)  # type: ignore[operator]
 
 
 @_power.register(types.CSBase | types.CupyCSMatrix)

src/fast_array_utils/stats/_typing.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: MPL-2.0
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal, Protocol, TypedDict
+from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypedDict, TypeVar
 
 import numpy as np
 
@@ -51,5 +51,8 @@ def __call__(
 Ops: TypeAlias = NoDtypeOps | DtypeOps
 
 
-class DTypeKw(TypedDict, total=False):
-    dtype: DTypeLike
+_DT = TypeVar("_DT", bound="DTypeLike")
+
+
+class DTypeKw(TypedDict, Generic[_DT], total=False):
+    dtype: _DT

src/fast_array_utils/stats/_utils.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import TYPE_CHECKING, Literal, cast, get_args
+from typing import TYPE_CHECKING, Literal, TypeVar, cast, get_args
 
 import numpy as np
 from numpy.exceptions import AxisError
@@ -71,7 +71,7 @@ def _dask_block(
     fns = {fn.__name__: fn for fn in (min, max, sum)}
 
     axis = _normalize_axis(axis, a.ndim)
-    rv = fns[op](a, axis=axis, keep_cupy_as_array=True, **_dtype_kw(dtype, op))  # type: ignore[misc,call-overload]
+    rv = fns[op](a, axis=axis, keep_cupy_as_array=True, **_dtype_kw(dtype, op))  # type: ignore[call-overload]
     shape = _get_shape(rv, axis=axis, keepdims=keepdims)
     return cast("NDArray[Any] | types.CupyArray", rv.reshape(shape))
 
@@ -109,5 +109,8 @@ def _get_shape(a: NDArray[Any] | np.number[Any] | types.CupyArray, *, axis: Lite
             raise AssertionError(msg)
 
 
-def _dtype_kw(dtype: DTypeLike | None, op: Ops) -> DTypeKw:
+DT = TypeVar("DT", bound="DTypeLike")
+
+
+def _dtype_kw(dtype: DT | None, op: Ops) -> DTypeKw[DT]:
     return {"dtype": dtype} if dtype is not None and op in get_args(DtypeOps) else {}