TYP: Type annotations overhaul, episode 2 (#288)

* TYP: annotate `_internal.get_xp` (and curse at `ParamSpec` for being so useless)

* TYP: fix (or ignore) typing errors in `common._helpers` (and curse at cupy)

* TYP: fix typing errors in `common._fft`

* TYP: fix typing errors in `common._aliases`

* TYP: fix typing errors in `common._linalg`

* TYP: fix/ignore typing errors in `numpy.__init__`

* TYP: fix typing errors in `numpy._typing`

* TYP: fix typing errors in `numpy._aliases`

* TYP: fix typing errors in `numpy._info`

* TYP: fix typing errors in `numpy._fft`

* TYP: it's a bad idea to import `TypeAlias` from `typing` on `python<3.10`

* TYP: it's also a bad idea to import `TypeGuard` from `typing` on `python<3.10`

* TYP: don't scare the prehistoric `dtype` from numpy 1.21

* TYP: dust off the DeLorean

* TYP: figure out how to drive a DeLorean

* TYP: apply review suggestions

Co-authored-by: crusaderky <[email protected]>

* TYP: sprinkle some `TypeAlias`es and `Final`s around

* TYP: `__dir__`

* TYP: fix typing errors in `numpy.linalg`

* TYP: add a `common._typing.Capabilities` typed dict type

* TYP: `__array_namespace_info__` helper types

* TYP: `dask.array` typing fixes and improvements

* STY: give the `=` some breathing room

Co-authored-by: Lucas Colley <[email protected]>

* STY: apply review suggestions

Co-authored-by: lucascolley <[email protected]>

---------

Co-authored-by: crusaderky <[email protected]>
Co-authored-by: Lucas Colley <[email protected]>

Author: 3 people

array_api_compat/_internal.py

@@ -2,10 +2,16 @@
 Internal helpers
 """
 
+from collections.abc import Callable
 from functools import wraps
 from inspect import signature
+from types import ModuleType
+from typing import TypeVar
 
-def get_xp(xp):
+_T = TypeVar("_T")
+
+
+def get_xp(xp: ModuleType) -> Callable[[Callable[..., _T]], Callable[..., _T]]:
     """
     Decorator to automatically replace xp with the corresponding array module.
 
@@ -22,14 +28,14 @@ def func(x, /, xp, kwarg=None):
 
     """
 
-    def inner(f):
+    def inner(f: Callable[..., _T], /) -> Callable[..., _T]:
         @wraps(f)
-        def wrapped_f(*args, **kwargs):
+        def wrapped_f(*args: object, **kwargs: object) -> object:
             return f(*args, xp=xp, **kwargs)
 
         sig = signature(f)
         new_sig = sig.replace(
-            parameters=[sig.parameters[i] for i in sig.parameters if i != "xp"]
+            parameters=[par for i, par in sig.parameters.items() if i != "xp"]
         )
 
         if wrapped_f.__doc__ is None:
@@ -40,7 +46,14 @@ def wrapped_f(*args, **kwargs):
 specification for more details.
 
 """
-        wrapped_f.__signature__ = new_sig
-        return wrapped_f
+        wrapped_f.__signature__ = new_sig  # pyright: ignore[reportAttributeAccessIssue]
+        return wrapped_f  # pyright: ignore[reportReturnType]
 
     return inner
+
+
+__all__ = ["get_xp"]
+
+
+def __dir__() -> list[str]:
+    return __all__

array_api_compat/common/__init__.py

@@ -1 +1 @@
-from ._helpers import * # noqa: F403
+from ._helpers import *  # noqa: F403