1
1
from __future__ import annotations
2
2
3
3
from types import ModuleType as Namespace
4
- from typing import Any , Protocol , TypeAlias , TypeVar
4
+ from typing import Any , Protocol , TypeAlias , TypedDict , TypeVar
5
5
6
6
_T_co = TypeVar ("_T_co" , covariant = True )
7
7
8
+
8
9
class NestedSequence (Protocol [_T_co ]):
9
10
def __getitem__ (self , key : int , / ) -> _T_co | NestedSequence [_T_co ]: ...
10
11
def __len__ (self , / ) -> int : ...
@@ -19,6 +20,16 @@ class HasShape(Protocol[_T_co]):
19
20
def shape (self , / ) -> _T_co : ...
20
21
21
22
23
+ Capabilities = TypedDict (
24
+ "Capabilities" ,
25
+ {
26
+ "boolean indexing" : bool ,
27
+ "data-dependent shapes" : bool ,
28
+ "max dimensions" : int ,
29
+ },
30
+ )
31
+
32
+
22
33
SupportsBufferProtocol : TypeAlias = Any
23
34
Array : TypeAlias = Any
24
35
Device : TypeAlias = Any
@@ -27,12 +38,13 @@ def shape(self, /) -> _T_co: ...
27
38
28
39
__all__ = [
29
40
"Array" ,
30
- "SupportsArrayNamespace " ,
41
+ "Capabilities " ,
31
42
"DType" ,
32
43
"Device" ,
33
44
"HasShape" ,
34
45
"Namespace" ,
35
46
"NestedSequence" ,
47
+ "SupportsArrayNamespace" ,
36
48
"SupportsBufferProtocol" ,
37
49
]
38
50
0 commit comments