diff options
Diffstat (limited to 'numpy/array_api/_typing.py')
| -rw-r--r-- | numpy/array_api/_typing.py | 30 |
1 files changed, 23 insertions, 7 deletions
diff --git a/numpy/array_api/_typing.py b/numpy/array_api/_typing.py index 5f937a56c..519e8463c 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, TypeVar -from . import ( - Array, +from ._array_object import Array +from numpy import ( + dtype, int8, int16, int32, @@ -33,12 +35,26 @@ from . import ( # This should really be recursive, but that isn't supported yet. See the # similar comment in numpy/typing/_array_like.py -NestedSequence = Sequence[Sequence[Any]] +_T = TypeVar("_T") +NestedSequence = Sequence[Sequence[_T]] 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 |