linters merge from main

Author: crusaderky

pixi.lock

@@ -61,7 +61,6 @@ basedpyright = "*"
 numpydoc = ">=1.8.0,<2"
 # import dependencies for mypy:
 array-api-strict = "*"
-numpy = "*"
 pytest = "*"
 
 [tool.pixi.feature.lint.tasks]

pyproject.toml

@@ -223,12 +223,12 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
 
     # Backend-specific branches
     if is_dask_namespace(xp):
-        import dask  # type: ignore[import-not-found]  # pylint: disable=import-outside-toplevel,import-error  # pyright: ignore[reportMissingImports]
+        import dask  # pylint: disable=import-outside-toplevel
 
         metas = [arg._meta for arg in args if hasattr(arg, "_meta")]  # pylint: disable=protected-access
         meta_xp = array_namespace(*metas)
 
-        wrapped = dask.delayed(
+        wrapped = dask.delayed(  # type: ignore[attr-defined]  # pyright: ignore[reportPrivateImportUsage]
             _lazy_apply_wrapper(func, as_numpy, multi_output, meta_xp),
             pure=True,
         )
@@ -239,7 +239,7 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
 
         out = tuple(
             xp.from_delayed(
-                delayed_out[i],
+                delayed_out[i],  # pyright: ignore[reportIndexIssue]
                 # Dask's unknown shapes diverge from the Array API specification
                 shape=tuple(math.nan if s is None else s for s in shape),
                 dtype=dtype,
@@ -254,7 +254,7 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         # Instead, we delay calling wrapped, which will receive
         # as arguments and will return JAX eager arrays.
 
-        import jax  # type: ignore[import-not-found]  # pylint: disable=import-outside-toplevel,import-error  # pyright: ignore[reportMissingImports]
+        import jax  # pylint: disable=import-outside-toplevel
 
         wrapped = _lazy_apply_wrapper(func, as_numpy, multi_output, xp)
 
@@ -265,18 +265,17 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
             out = wrapped(*args, **kwargs)
 
         else:
-            out = cast(
-                tuple[Array, ...],
-                jax.pure_callback(
-                    wrapped,
-                    tuple(
-                        jax.ShapeDtypeStruct(shape, dtype)  # pyright: ignore[reportUnknownArgumentType]
-                        for shape, dtype in zip(shapes, dtypes, strict=True)
-                    ),
-                    *args,
-                    **kwargs,
+            # FIXME jax typing bug
+            out_jax = jax.pure_callback(  # type: ignore[func-returns-value]
+                wrapped,
+                tuple(
+                    jax.ShapeDtypeStruct(shape, dtype)
+                    for shape, dtype in zip(shapes, dtypes, strict=True)
                 ),
+                *args,
+                **kwargs,
             )
+            out = cast(tuple[Array, ...], cast(object, out_jax))
 
     else:
         # Eager backends

src/array_api_extra/_lib/_lazy.py

@@ -18,7 +18,7 @@
         size,
     )
 except ImportError:
-    from array_api_compat import (  # pyright: ignore[reportMissingTypeStubs]
+    from array_api_compat import (  # pyright: ignore[reportMissingImports]
         array_namespace,
         device,
         is_array_api_strict_namespace,

src/array_api_extra/_lib/_utils/_compat.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 import pytest
-from array_api_compat import (  # type: ignore[import-untyped]  # pyright: ignore[reportMissingTypeStubs]
+from array_api_compat import (  # type: ignore[import-not-found]  # pyright: ignore[reportMissingImports]
     array_namespace,
     is_writeable_array,
 )