Skip to content

Commit 9acba46

Browse files
committed
TYP: add a common._typing.Capabilities typed dict type
1 parent c66b750 commit 9acba46

File tree

1 file changed

+14
-2
lines changed

1 file changed

+14
-2
lines changed

array_api_compat/common/_typing.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
from __future__ import annotations
22

33
from types import ModuleType as Namespace
4-
from typing import Any, Protocol, TypeAlias, TypeVar
4+
from typing import Any, Protocol, TypeAlias, TypedDict, TypeVar
55

66
_T_co = TypeVar("_T_co", covariant=True)
77

8+
89
class NestedSequence(Protocol[_T_co]):
910
def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
1011
def __len__(self, /) -> int: ...
@@ -19,6 +20,16 @@ class HasShape(Protocol[_T_co]):
1920
def shape(self, /) -> _T_co: ...
2021

2122

23+
Capabilities = TypedDict(
24+
"Capabilities",
25+
{
26+
"boolean indexing": bool,
27+
"data-dependent shapes": bool,
28+
"max dimensions": int,
29+
},
30+
)
31+
32+
2233
SupportsBufferProtocol: TypeAlias = Any
2334
Array: TypeAlias = Any
2435
Device: TypeAlias = Any
@@ -27,12 +38,13 @@ def shape(self, /) -> _T_co: ...
2738

2839
__all__ = [
2940
"Array",
30-
"SupportsArrayNamespace",
41+
"Capabilities",
3142
"DType",
3243
"Device",
3344
"HasShape",
3445
"Namespace",
3546
"NestedSequence",
47+
"SupportsArrayNamespace",
3648
"SupportsBufferProtocol",
3749
]
3850

0 commit comments

Comments
 (0)