diff options
| -rw-r--r-- | numpy/array_api/_typing.py | 27 |
1 files changed, 21 insertions, 6 deletions
diff --git a/numpy/array_api/_typing.py b/numpy/array_api/_typing.py index 5f937a56c..f66279fbc 100644 --- a/numpy/array_api/_typing.py +++ b/numpy/array_api/_typing.py @@ -15,10 +15,12 @@ __all__ = [ "PyCapsule", ] -from typing import Any, Literal, Sequence, Type, Union +import sys +from typing import Any, Literal, Sequence, Type, Union, TYPE_CHECKING -from . import ( - Array, +from . import Array +from numpy import ( + dtype, int8, int16, int32, @@ -36,9 +38,22 @@ from . import ( NestedSequence = Sequence[Sequence[Any]] Device = Literal["cpu"] -Dtype = Type[ - Union[int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64] -] +if TYPE_CHECKING or sys.version_info >= (3, 9): + Dtype = dtype[Union[ + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, + ]] +else: + Dtype = dtype + SupportsDLPack = Any SupportsBufferProtocol = Any PyCapsule = Any |