TYP: sprinkle some TypeAliases and Finals around

Author: jorenham

array_api_compat/common/_fft.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Literal
+from typing import Literal, TypeAlias
 
 from ._typing import Array, Device, DType, Namespace
 
-_Norm = Literal["backward", "ortho", "forward"]
+_Norm: TypeAlias = Literal["backward", "ortho", "forward"]
 
 # Note: NumPy fft functions improperly upcast float32 and complex64 to
 # complex128, which is why we require wrapping them all here.

array_api_compat/common/_typing.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from types import ModuleType as Namespace
-from typing import Any, Protocol, TypeVar
+from typing import Any, Protocol, TypeAlias, TypeVar
 
 __all__ = [
     "Array",
@@ -30,7 +30,7 @@ class HasShape(Protocol[_T_co]):
     def shape(self, /) -> _T_co: ...
 
 
-SupportsBufferProtocol = Any
-Array = Any
-Device = Any
-DType = Any
+SupportsBufferProtocol: TypeAlias = Any
+Array: TypeAlias = Any
+Device: TypeAlias = Any
+DType: TypeAlias = Any

array_api_compat/numpy/__init__.py

@@ -32,4 +32,4 @@
 except ImportError:
     pass
 
-__array_api_version__ = "2024.12"
+__array_api_version__: Final = "2024.12"

array_api_compat/numpy/_typing.py

@@ -3,13 +3,13 @@
 __all__ = ["Array", "DType", "Device"]
 _all_ignore = ["np"]
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, TypeAlias
 
 import numpy as np
 
-Device = Literal["cpu"]
+Device: TypeAlias = Literal["cpu"]
+
 if TYPE_CHECKING:
-    from typing_extensions import TypeAlias
 
     # NumPy 1.x on Python 3.10 fails to parse np.dtype[]
     DType: TypeAlias = np.dtype[
@@ -22,5 +22,5 @@
     ]
     Array: TypeAlias = np.ndarray[Any, DType]
 else:
-    DType = np.dtype
-    Array = np.ndarray
+    DType: TypeAlias = np.dtype
+    Array: TypeAlias = np.ndarray