✨ HasSize

Signed-off-by: Nathaniel Starkman <nstarman@users.noreply.github.com>

Author: nstarman

src/array_api_typing/__init__.py

@@ -7,6 +7,7 @@
     "HasMatrixTranspose",
     "HasNDim",
     "HasShape",
+    "HasSize",
     "__version__",
     "__version_tuple__",
 )
@@ -18,5 +19,6 @@
     HasMatrixTranspose,
     HasNDim,
     HasShape,
+    HasSize,
 )
 from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_array.py

@@ -6,6 +6,7 @@
     "HasMatrixTranspose",
     "HasNDim",
     "HasShape",
+    "HasSize",
 )
 
 from types import ModuleType
@@ -137,13 +138,33 @@ def shape(self) -> tuple[int | None, ...]:
         ...
 
 
+class HasSize(Protocol):
+    """Protocol for array classes that have a size attribute."""
+
+    @property
+    def size(self) -> int | None:
+        """Number of elements in an array.
+
+        Returns:
+            int | None: number of elements in an array. The returned value must
+                be `None` if and only if one or more array dimensions are
+                unknown.
+
+        Notes:
+            This must equal the product of the array's dimensions.
+
+        """
+        ...
+
+
 class Array(
     # ------ Attributes -------
     HasDType[DTypeT_co],
     HasDevice,
     HasMatrixTranspose,
     HasNDim,
     HasShape,
+    HasSize,
     # ------- Methods ---------
     HasArrayNamespace[NamespaceT_co],
     # -------------------------

tests/integration/test_numpy1p0.pyi

@@ -71,3 +71,7 @@ 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, ...])
+
+# Check Attribute `.size`
+assert_type(x_f32.size, int | None)
+assert_type(x_i32.size, int | None)

tests/integration/test_numpy2p0.pyi

@@ -81,3 +81,8 @@ assert_type(x_b.ndim, int)
 assert_type(x_f32.shape, tuple[int | None, ...])
 assert_type(x_i32.shape, tuple[int | None, ...])
 assert_type(x_b.shape, tuple[int | None, ...])
+
+# Check Attribute `.size`
+assert_type(x_f32.size, int | None)
+assert_type(x_i32.size, int | None)
+assert_type(x_b.size, int | None)