✨ HasShape

Author: nstarman

src/array_api_typing/__init__.py

@@ -6,6 +6,7 @@
     "HasDType",
     "HasMatrixTranspose",
     "HasNDim",
+    "HasShape",
     "__version__",
     "__version_tuple__",
 )
@@ -16,5 +17,6 @@
     HasDType,
     HasMatrixTranspose,
     HasNDim,
+    HasShape,
 )
 from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_array.py

@@ -5,6 +5,7 @@
     "HasDevice",
     "HasMatrixTranspose",
     "HasNDim",
+    "HasShape",
 )
 
 from types import ModuleType
@@ -115,12 +116,34 @@ def ndim(self) -> int:
         ...
 
 
+class HasShape(Protocol):
+    """Protocol for array classes that have a shape attribute."""
+
+    @property
+    def shape(self) -> tuple[int | None, ...]:
+        """Shape of the array.
+
+        Returns:
+            tuple[int | None, ...]: array dimensions. An array dimension must be None
+                if and only if a dimension is unknown.
+
+        Notes:
+            For array libraries having graph-based computational models, array
+            dimensions may be unknown due to data-dependent operations (e.g.,
+            boolean indexing; `A[:, B > 0]`) and thus cannot be statically
+            resolved without knowing array contents.
+
+        """
+        ...
+
+
 class Array(
     # ------ Attributes -------
     HasDType[DTypeT_co],
     HasDevice,
     HasMatrixTranspose,
     HasNDim,
+    HasShape,
     # ------- Methods ---------
     HasArrayNamespace[NamespaceT_co],
     # -------------------------

tests/integration/test_numpy1p0.pyi

@@ -67,3 +67,7 @@ assert_type(x_i32.mT, xpt.Array[dtype[Any]])
 # Check Attribute `.ndim`
 assert_type(x_f32.ndim, int)
 assert_type(x_i32.ndim, int)
+
+# Check Attribute `.shape`
+assert_type(x_f32.shape, tuple[int | None, ...])
+assert_type(x_i32.shape, tuple[int | None, ...])

tests/integration/test_numpy2p0.pyi

@@ -76,3 +76,8 @@ assert_type(x_b.mT, xpt.Array[np.dtype[B]])
 assert_type(x_f32.ndim, int)
 assert_type(x_i32.ndim, int)
 assert_type(x_b.ndim, int)
+
+# Check Attribute `.shape`
+assert_type(x_f32.shape, tuple[int | None, ...])
+assert_type(x_i32.shape, tuple[int | None, ...])
+assert_type(x_b.shape, tuple[int | None, ...])