TYP: apply review suggestions

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

Author: jorenham

array_api_compat/_internal.py

@@ -2,21 +2,18 @@
 Internal helpers
 """
 
+from collections.abc import Callable
 from functools import wraps
 from inspect import signature
-from typing import TYPE_CHECKING
+from types import ModuleType
+from typing import TypeVar
 
 __all__ = ["get_xp"]
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
-    from types import ModuleType
-    from typing import TypeVar
+_T = TypeVar("_T")
 
-    _T = TypeVar("_T")
 
-
-def get_xp(xp: "ModuleType") -> "Callable[[Callable[..., _T]], Callable[..., _T]]":
+def get_xp(xp: ModuleType) -> Callable[[Callable[..., _T]], Callable[..., _T]]:
     """
     Decorator to automatically replace xp with the corresponding array module.
 
@@ -33,7 +30,7 @@ def func(x, /, xp, kwarg=None):
 
     """
 
-    def inner(f: "Callable[..., _T]", /) -> "Callable[..., _T]":
+    def inner(f: Callable[..., _T], /) -> Callable[..., _T]:
         @wraps(f)
         def wrapped_f(*args: object, **kwargs: object) -> object:
             return f(*args, xp=xp, **kwargs)

array_api_compat/common/_aliases.py

@@ -13,6 +13,7 @@
 from ._typing import Array, Device, DType, Namespace
 
 if TYPE_CHECKING:
+    # TODO: import from typing (requires Python >=3.13)
     from typing_extensions import TypeIs
 
 # These functions are modified from the NumPy versions.

array_api_compat/common/_helpers.py

@@ -12,12 +12,22 @@
 import math
 import sys
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, SupportsIndex, cast, overload
+from collections.abc import Collection
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    SupportsIndex,
+    TypeAlias,
+    TypeGuard,
+    TypeVar,
+    cast,
+    overload,
+)
 
 from ._typing import Array, Device, HasShape, Namespace, SupportsArrayNamespace
 
 if TYPE_CHECKING:
-    from collections.abc import Collection
 
     import dask.array as da
     import jax
@@ -26,7 +36,9 @@
     import numpy.typing as npt
     import sparse  # pyright: ignore[reportMissingTypeStubs]
     import torch
-    from typing_extensions import TypeAlias, TypeGuard, TypeIs, TypeVar
+
+    # TODO: import from typing (requires Python >=3.13)
+    from typing_extensions import TypeIs, TypeVar
 
     _SizeT = TypeVar("_SizeT", bound=int | None)
 

array_api_compat/numpy/_aliases.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from builtins import bool as py_bool
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast
 
 import numpy as np
 
@@ -13,11 +13,9 @@
 from ._typing import Array, Device, DType
 
 if TYPE_CHECKING:
-    from typing import Any, Literal
+    from typing_extensions import Buffer, TypeIs
 
-    from typing_extensions import Buffer, TypeAlias, TypeIs
-
-    _Copy: TypeAlias = py_bool | Literal[2] | np._CopyMode
+_Copy: TypeAlias = py_bool | Literal[2] | np._CopyMode
 
 bool = np.bool_