fix mypy errors

Author: lucascolley

src/array_api_extra/_funcs.py

@@ -5,18 +5,17 @@
 
 import operator
 import warnings
-
-# https://github.com/pylint-dev/pylint/issues/10112
-from collections.abc import Callable  # pylint: disable=import-error
-from typing import ClassVar, Literal
+from collections.abc import Callable
+from types import ModuleType
+from typing import ClassVar, Literal, cast
 
 from ._lib import _compat, _utils
 from ._lib._compat import (
     array_namespace,
     is_jax_array,
     is_writeable_array,
 )
-from ._lib._typing import Array, Index, ModuleType
+from ._lib._typing import Array, Index
 
 __all__ = [
     "at",
@@ -779,7 +778,7 @@ def _update_common(
         if copy:
             if is_jax_array(x):
                 # Use JAX's at[]
-                func = getattr(x.at[idx], at_op)
+                func = cast(Callable[[Array], Array], getattr(x.at[idx], at_op))
                 return func(y), None
             # Emulate at[] behaviour for non-JAX arrays
             # with a copy followed by an update

tests/test_at.py

@@ -1,6 +1,7 @@
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from importlib import import_module
+from typing import cast
 
 import numpy as np
 import pytest
@@ -100,7 +101,8 @@ def test_update_ops(
         pytest.skip("at() does not support updates on sparse arrays")
 
     with assert_copy(array, expect_copy):
-        y = getattr(at(array)[1:], op)(arg, **kwargs)
+        func = cast(Callable[..., Array], getattr(at(array)[1:], op))  # type: ignore[no-any-explicit]
+        y = func(arg, **kwargs)
         assert isinstance(y, type(array))
         assert_array_equal(y, expect)
 
@@ -153,6 +155,6 @@ def test_iops_incompatible_dtype(op: str, copy: bool):
     to dtype('int64') with casting rule 'same_kind'
     """
     a = np.asarray([2, 4])
-    func = getattr(at(a)[:], op)
+    func = cast(Callable[..., Array], getattr(at(a)[:], op))  # type: ignore[no-any-explicit]
     with pytest.raises(TypeError, match="Cannot cast ufunc"):
         func(1.1, copy=copy)