TYP: dask.array typing fixes and improvements

Author: jorenham

array_api_compat/dask/array/__init__.py

@@ -1,9 +1,11 @@
-from dask.array import * # noqa: F403
+from typing import Final
+
+from dask.array import *  # noqa: F403
 
 # These imports may overwrite names from the import * above.
-from ._aliases import * # noqa: F403
+from ._aliases import *  # noqa: F403
 
-__array_api_version__ = '2024.12'
+__array_api_version__: Final = "2024.12"
 
 # See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')

array_api_compat/dask/array/_aliases.py

@@ -1,28 +1,38 @@
+# pyright: reportPrivateUsage=false
+# pyright: reportUnknownArgumentType=false
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false
+
 from __future__ import annotations
 
-from typing import Callable, Optional, Union
+from builtins import bool as py_bool
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from typing_extensions import TypeIs
 
+import dask.array as da
 import numpy as np
+from numpy import bool_ as bool
 from numpy import (
-    # dtypes
-    bool_ as bool,
+    can_cast,
+    complex64,
+    complex128,
     float32,
     float64,
     int8,
     int16,
     int32,
     int64,
+    result_type,
     uint8,
     uint16,
     uint32,
     uint64,
-    complex64,
-    complex128,
-    can_cast,
-    result_type,
 )
-import dask.array as da
 
+from ..._internal import get_xp
 from ...common import _aliases, _helpers, array_namespace
 from ...common._typing import (
     Array,
@@ -31,7 +41,6 @@
     NestedSequence,
     SupportsBufferProtocol,
 )
-from ..._internal import get_xp
 from ._info import __array_namespace_info__
 
 isdtype = get_xp(np)(_aliases.isdtype)
@@ -44,8 +53,8 @@ def astype(
     dtype: DType,
     /,
     *,
-    copy: bool = True,
-    device: Optional[Device] = None,
+    copy: py_bool = True,
+    device: Device | None = None,
 ) -> Array:
     """
     Array API compatibility wrapper for astype().
@@ -69,14 +78,14 @@ def astype(
 # not pass stop/step as keyword arguments, which will cause
 # an error with dask
 def arange(
-    start: Union[int, float],
+    start: float,
     /,
-    stop: Optional[Union[int, float]] = None,
-    step: Union[int, float] = 1,
+    stop: float | None = None,
+    step: float = 1,
     *,
-    dtype: Optional[DType] = None,
-    device: Optional[Device] = None,
-    **kwargs,
+    dtype: DType | None = None,
+    device: Device | None = None,
+    **kwargs: object,
 ) -> Array:
     """
     Array API compatibility wrapper for arange().
@@ -87,7 +96,7 @@ def arange(
     # TODO: respect device keyword?
     _helpers._check_device(da, device)
 
-    args = [start]
+    args: list[Any] = [start]
     if stop is not None:
         args.append(stop)
     else:
@@ -137,18 +146,13 @@ def arange(
 
 # asarray also adds the copy keyword, which is not present in numpy 1.0.
 def asarray(
-    obj: (
-        Array 
-        | bool | int | float | complex 
-        | NestedSequence[bool | int | float | complex] 
-        | SupportsBufferProtocol
-    ),
+    obj: complex | NestedSequence[complex] | Array | SupportsBufferProtocol,
     /,
     *,
-    dtype: Optional[DType] = None,
-    device: Optional[Device] = None,
-    copy: Optional[bool] = None,
-    **kwargs,
+    dtype: DType | None = None,
+    device: Device | None = None,
+    copy: py_bool | None = None,
+    **kwargs: object,
 ) -> Array:
     """
     Array API compatibility wrapper for asarray().
@@ -164,7 +168,7 @@ def asarray(
             if copy is False:
                 raise ValueError("Unable to avoid copy when changing dtype")
             obj = obj.astype(dtype)
-        return obj.copy() if copy else obj
+        return obj.copy() if copy else obj  # pyright: ignore[reportAttributeAccessIssue]
 
     if copy is False:
         raise NotImplementedError(
@@ -177,22 +181,21 @@ def asarray(
     return da.from_array(obj)
 
 
-from dask.array import (
-    # Element wise aliases
-    arccos as acos,
-    arccosh as acosh,
-    arcsin as asin,
-    arcsinh as asinh,
-    arctan as atan,
-    arctan2 as atan2,
-    arctanh as atanh,
-    left_shift as bitwise_left_shift,
-    right_shift as bitwise_right_shift,
-    invert as bitwise_invert,
-    power as pow,
-    # Other
-    concatenate as concat,
-)
+# Element wise aliases
+from dask.array import arccos as acos
+from dask.array import arccosh as acosh
+from dask.array import arcsin as asin
+from dask.array import arcsinh as asinh
+from dask.array import arctan as atan
+from dask.array import arctan2 as atan2
+from dask.array import arctanh as atanh
+
+# Other
+from dask.array import concatenate as concat
+from dask.array import invert as bitwise_invert
+from dask.array import left_shift as bitwise_left_shift
+from dask.array import power as pow
+from dask.array import right_shift as bitwise_right_shift
 
 
 # dask.array.clip does not work unless all three arguments are provided.
@@ -202,8 +205,8 @@ def asarray(
 def clip(
     x: Array,
     /,
-    min: Optional[Union[int, float, Array]] = None,
-    max: Optional[Union[int, float, Array]] = None,
+    min: float | Array | None = None,
+    max: float | Array | None = None,
 ) -> Array:
     """
     Array API compatibility wrapper for clip().
@@ -212,8 +215,8 @@ def clip(
     specification for more details.
     """
 
-    def _isscalar(a):
-        return isinstance(a, (int, float, type(None)))
+    def _isscalar(a: float | Array | None, /) -> TypeIs[float | None]:
+        return a is None or isinstance(a, (int, float))
 
     min_shape = () if _isscalar(min) else min.shape
     max_shape = () if _isscalar(max) else max.shape
@@ -266,7 +269,12 @@ def _ensure_single_chunk(x: Array, axis: int) -> tuple[Array, Callable[[Array],
 
 
 def sort(
-    x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True
+    x: Array,
+    /,
+    *,
+    axis: int = -1,
+    descending: py_bool = False,
+    stable: py_bool = True,
 ) -> Array:
     """
     Array API compatibility layer around the lack of sort() in Dask.
@@ -296,7 +304,12 @@ def sort(
 
 
 def argsort(
-    x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True
+    x: Array,
+    /,
+    *,
+    axis: int = -1,
+    descending: py_bool = False,
+    stable: py_bool = True,
 ) -> Array:
     """
     Array API compatibility layer around the lack of argsort() in Dask.
@@ -330,25 +343,34 @@ def argsort(
 # dask.array.count_nonzero does not have keepdims
 def count_nonzero(
     x: Array,
-    axis=None,
-    keepdims=False
+    axis: int | None = None,
+    keepdims: py_bool = False,
 ) -> Array:
-   result = da.count_nonzero(x, axis)
-   if keepdims:
-       if axis is None:
-            return da.reshape(result, [1]*x.ndim)
-       return da.expand_dims(result, axis)
-   return result
-
-
+    result = da.count_nonzero(x, axis)
+    if keepdims:
+        if axis is None:
+            return da.reshape(result, [1] * x.ndim)
+        return da.expand_dims(result, axis)
+    return result
+
+
+__all__ = [
+    "__array_namespace_info__",
+    "count_nonzero",
+    "bool",
+    "int8", "int16", "int32", "int64",
+    "uint8", "uint16", "uint32", "uint64",
+    "float32", "float64",
+    "complex64", "complex128",
+    "asarray", "astype", "can_cast", "result_type",
+    "pow",
+    "concat",
+    "acos", "acosh", "asin", "asinh", "atan", "atan2", "atanh",
+    "bitwise_left_shift", "bitwise_right_shift", "bitwise_invert",
+]  # fmt: skip
+__all__ += _aliases.__all__
+_all_ignore = ["array_namespace", "get_xp", "da", "np"]
 
-__all__ = _aliases.__all__ + [
-                    '__array_namespace_info__', 'asarray', 'astype', 'acos',
-                    'acosh', 'asin', 'asinh', 'atan', 'atan2',
-                    'atanh', 'bitwise_left_shift', 'bitwise_invert',
-                    'bitwise_right_shift', 'concat', 'pow', 'can_cast',
-                    'result_type', 'bool', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64',
-                    'uint8', 'uint16', 'uint32', 'uint64', 'complex64', 'complex128',
-                    'can_cast', 'count_nonzero', 'result_type']
 
-_all_ignore = ["array_namespace", "get_xp", "da", "np"]
+def __dir__() -> list[str]:
+    return __all__