3
3
import sys
4
4
from typing import Any , overload , Sequence , TYPE_CHECKING , Union , TypeVar
5
5
6
- from numpy import ndarray , dtype
7
- from ._scalars import _ScalarLike
6
+ from numpy import (
7
+ ndarray ,
8
+ dtype ,
9
+ generic ,
10
+ bool_ ,
11
+ unsignedinteger ,
12
+ integer ,
13
+ floating ,
14
+ complexfloating ,
15
+ timedelta64 ,
16
+ datetime64 ,
17
+ object_ ,
18
+ void ,
19
+ str_ ,
20
+ bytes_ ,
21
+ )
8
22
from ._dtype_like import DTypeLike
9
23
10
24
if sys .version_info >= (3 , 8 ):
18
32
else :
19
33
HAVE_PROTOCOL = True
20
34
35
+ _T = TypeVar ("_T" )
21
36
_DType = TypeVar ("_DType" , bound = "dtype[Any]" )
22
37
23
38
if TYPE_CHECKING or HAVE_PROTOCOL :
@@ -30,6 +45,24 @@ def __array__(self, dtype: None = ...) -> ndarray[Any, _DType]: ...
30
45
else :
31
46
_SupportsArray = Any
32
47
48
+ # TODO: Wait for support for recursive types
49
+ _NestedSequence = Union [
50
+ _T ,
51
+ Sequence [_T ],
52
+ Sequence [Sequence [_T ]],
53
+ Sequence [Sequence [Sequence [_T ]]],
54
+ Sequence [Sequence [Sequence [Sequence [_T ]]]],
55
+ ]
56
+ _RecursiveSequence = Sequence [Sequence [Sequence [Sequence [Sequence [Any ]]]]]
57
+
58
+ # A union representing array-like objects; consists of two typevars:
59
+ # One representing types that can be parametrized w.r.t. `np.dtype`
60
+ # and another one for the rest
61
+ _ArrayLike = Union [
62
+ _NestedSequence [_SupportsArray [_DType ]],
63
+ _NestedSequence [_T ],
64
+ ]
65
+
33
66
# TODO: support buffer protocols once
34
67
#
35
68
# https://bugs.python.org/issue27501
@@ -38,8 +71,48 @@ def __array__(self, dtype: None = ...) -> ndarray[Any, _DType]: ...
38
71
#
39
72
# https://github.com/python/typing/issues/593
40
73
ArrayLike = Union [
41
- _ScalarLike ,
42
- Sequence [_ScalarLike ],
43
- Sequence [Sequence [Any ]], # TODO: Wait for support for recursive types
44
- "_SupportsArray[Any]" ,
74
+ _RecursiveSequence ,
75
+ _ArrayLike [
76
+ "dtype[Any]" ,
77
+ Union [bool , int , float , complex , str , bytes ]
78
+ ],
79
+ ]
80
+
81
+ # `ArrayLike<X>`: array-like objects that can be coerced into `X`
82
+ # given the casting rules `same_kind`
83
+ _ArrayLikeBool = _ArrayLike [
84
+ "dtype[bool_]" ,
85
+ bool ,
86
+ ]
87
+ _ArrayLikeUInt = _ArrayLike [
88
+ "dtype[Union[bool_, unsignedinteger[Any]]]" ,
89
+ bool ,
90
+ ]
91
+ _ArrayLikeInt = _ArrayLike [
92
+ "dtype[Union[bool_, integer[Any]]]" ,
93
+ Union [bool , int ],
94
+ ]
95
+ _ArrayLikeFloat = _ArrayLike [
96
+ "dtype[Union[bool_, integer[Any], floating[Any]]]" ,
97
+ Union [bool , int , float ],
98
+ ]
99
+ _ArrayLikeComplex = _ArrayLike [
100
+ "dtype[Union[bool_, integer[Any], floating[Any], complexfloating[Any, Any]]]" ,
101
+ Union [bool , int , float , complex ],
102
+ ]
103
+ _ArrayLikeTD64 = _ArrayLike [
104
+ "dtype[Union[bool_, integer[Any], timedelta64]]" ,
105
+ Union [bool , int ],
106
+ ]
107
+ _ArrayLikeDT64 = _NestedSequence [_SupportsArray ["dtype[datetime64]" ]]
108
+ _ArrayLikeObject = _NestedSequence [_SupportsArray ["dtype[object_]" ]]
109
+
110
+ _ArrayLikeVoid = _NestedSequence [_SupportsArray ["dtype[void]" ]]
111
+ _ArrayLikeStr = _ArrayLike [
112
+ "dtype[str_]" ,
113
+ str ,
114
+ ]
115
+ _ArrayLikeBytes = _ArrayLike [
116
+ "dtype[bytes_]" ,
117
+ bytes ,
45
118
]
0 commit comments