data-apis · lucascolley · Apr 17, 2025 · Mar 22, 2025 · Mar 22, 2025 · Mar 22, 2025
diff --git a/array_api_compat/_internal.py b/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__
diff --git a/array_api_compat/common/__init__.py b/array_api_compat/common/__init__.py
@@ -1 +1 @@
-from ._helpers import * # noqa: F403
+from ._helpers import *  # noqa: F403
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		from ._helpers import * # noqa: F403
		from ._helpers import * # noqa: F403