TYP: Type annotations overhaul, episode 2

array_api_compat/_internal.py

@@ -4,8 +4,19 @@
 
 from functools import wraps
 from inspect import signature
+from typing import TYPE_CHECKING
 
-def get_xp(xp):
+__all__ = ["get_xp"]
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import ModuleType
+    from typing import TypeVar
+
+    _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 +33,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 +51,7 @@ 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