Skip to content

Commit e49a73a

Browse files
committed
✨ HasShape
1 parent ed33843 commit e49a73a

File tree

4 files changed

+34
-0
lines changed

4 files changed

+34
-0
lines changed

src/array_api_typing/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
"HasDevice",
88
"HasMatrixTranspose",
99
"HasNDim",
10+
"HasShape",
1011
"__version__",
1112
"__version_tuple__",
1213
)
@@ -18,5 +19,6 @@
1819
HasDType,
1920
HasMatrixTranspose,
2021
HasNDim,
22+
HasShape,
2123
)
2224
from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_array.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
"HasDevice",
66
"HasMatrixTranspose",
77
"HasNDim",
8+
"HasShape",
89
)
910

1011
from types import ModuleType
@@ -115,12 +116,34 @@ def ndim(self) -> int:
115116
...
116117

117118

119+
class HasShape(Protocol):
120+
"""Protocol for array classes that have a shape attribute."""
121+
122+
@property
123+
def shape(self) -> tuple[int | None, ...]:
124+
"""Shape of the array.
125+
126+
Returns:
127+
tuple[int | None, ...]: array dimensions. An array dimension must be None
128+
if and only if a dimension is unknown.
129+
130+
Notes:
131+
For array libraries having graph-based computational models, array
132+
dimensions may be unknown due to data-dependent operations (e.g.,
133+
boolean indexing; `A[:, B > 0]`) and thus cannot be statically
134+
resolved without knowing array contents.
135+
136+
"""
137+
...
138+
139+
118140
class Array(
119141
# ------ Attributes -------
120142
HasDType[DTypeT_co],
121143
HasDevice,
122144
HasMatrixTranspose,
123145
HasNDim,
146+
HasShape,
124147
# ------- Methods ---------
125148
HasArrayNamespace[NamespaceT_co],
126149
# -------------------------

tests/integration/test_numpy1p0.pyi

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,3 +67,7 @@ _: xpt.Array[dtype[Any]] = x_i32.mT
6767
# Check Attribute `.ndim`
6868
_: int = x_f32.ndim
6969
_: int = x_i32.ndim
70+
71+
# Check Attribute `.shape`
72+
_: tuple[int | None, ...] = x_f32.shape
73+
_: tuple[int | None, ...] = x_i32.shape

tests/integration/test_numpy2p0.pyi

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,3 +76,8 @@ _: xpt.Array[np.dtype[B]] = x_b.mT
7676
_: int = x_f32.ndim
7777
_: int = x_i32.ndim
7878
_: int = x_b.ndim
79+
80+
# Check Attribute `.shape`
81+
_: tuple[int | None, ...] = x_f32.shape
82+
_: tuple[int | None, ...] = x_i32.shape
83+
_: tuple[int | None, ...] = x_b.shape

0 commit comments

Comments
 (0)