lint

Author: crusaderky

src/array_api_extra/_lib/_funcs.py

@@ -126,22 +126,21 @@ def apply_where(  # type: ignore[no-any-explicit] # numpydoc ignore=PR01,PR02
     if (f2 is None) == (fill_value is None):
         msg = "Exactly one of `fill_value` or `f2` must be given."
         raise TypeError(msg)
-    if not isinstance(args, tuple):
-        args = (args,)
-    args = cast(tuple[Array, ...], args)
+    args_ = list(args) if isinstance(args, tuple) else [args]
+    del args
 
-    xp = array_namespace(cond, *args) if xp is None else xp
+    xp = array_namespace(cond, *args_) if xp is None else xp
 
     if getattr(fill_value, "ndim", 0):
-        cond, fill_value, *args = xp.broadcast_arrays(cond, fill_value, *args)
+        cond, fill_value, *args_ = xp.broadcast_arrays(cond, fill_value, *args_)
     else:
-        cond, *args = xp.broadcast_arrays(cond, *args)
+        cond, *args_ = xp.broadcast_arrays(cond, *args_)
 
     if is_dask_namespace(xp):
-        meta_xp = meta_namespace(cond, *args, fill_value, xp=xp)
+        meta_xp = meta_namespace(cond, fill_value, *args_, xp=xp)
         # map_blocks doesn't descend into tuples of Arrays
-        return xp.map_blocks(_apply_where, cond, f1, f2, fill_value, *args, xp=meta_xp)
-    return _apply_where(cond, f1, f2, fill_value, *args, xp=xp)
+        return xp.map_blocks(_apply_where, cond, f1, f2, fill_value, *args_, xp=meta_xp)
+    return _apply_where(cond, f1, f2, fill_value, *args_, xp=xp)
 
 
 def _apply_where(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR01,RT01

src/array_api_extra/_lib/_utils/_helpers.py

@@ -264,5 +264,5 @@ def meta_namespace(
     if not is_dask_namespace(xp):
         return xp
     # Quietly skip scalars and None's
-    metas = [getattr(a, "_meta", None) for a in arrays]
+    metas = [cast(Array | None, getattr(a, "_meta", None)) for a in arrays]
     return array_namespace(*metas)

tests/test_funcs.py

@@ -2,6 +2,7 @@
 import math
 import warnings
 from types import ModuleType
+from typing import Any, cast
 
 import hypothesis
 import hypothesis.extra.numpy as npst
@@ -29,7 +30,7 @@
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._helpers import asarrays, eager_shape, ndindex
-from array_api_extra._lib._utils._typing import Array, Device, DType
+from array_api_extra._lib._utils._typing import Array, Device
 from array_api_extra.testing import lazy_xp_function
 
 # some xp backends are untyped
@@ -120,7 +121,7 @@ def test_dtype_propagation(self, xp: ModuleType, library: Backend):
             cond,
             (x, y),
             self.f1,
-            lambda x, y: mxp.astype(x - y, xp.int64),  # pyright: ignore[reportUnknownArgumentType]
+            lambda x, y: mxp.astype(x - y, xp.int64),  # pyright: ignore[reportArgumentType,reportUnknownArgumentType]
         )
         assert actual.dtype == xp.int64
 
@@ -212,11 +213,11 @@ def test_device(self, xp: ModuleType, device: Device):
         p=st.floats(min_value=0, max_value=1),
         data=st.data(),
     )
-    def test_hypothesis(  # type: ignore[no-any-decorated]
+    def test_hypothesis(  # type: ignore[no-any-explicit,no-any-decorated]
         self,
         n_arrays: int,
         rng_seed: int,
-        dtype: DType,
+        dtype: np.dtype[Any],
         p: float,
         data: st.DataObject,
         xp: ModuleType,
@@ -242,10 +243,10 @@ def test_hypothesis(  # type: ignore[no-any-decorated]
         )
 
         def f1(*args: Array) -> Array:
-            return sum(args)
+            return cast(Array, sum(args))
 
         def f2(*args: Array) -> Array:
-            return sum(args) / 2
+            return cast(Array, sum(args) / 2)
 
         rng = np.random.default_rng(rng_seed)
         cond = xp.asarray(rng.random(size=cond_shape) > p)

tests/test_helpers.py

@@ -20,7 +20,7 @@
 
 # mypy: disable-error-code=no-untyped-usage
 
-np_compat = array_namespace(np.empty(0))
+np_compat = array_namespace(np.empty(0))  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
 # FIXME calls xp.unique_values without size
 lazy_xp_function(in1d, jax_jit=False, static_argnames=("assume_unique", "invert", "xp"))