types-array-api 1.1.1__tar.gz → 1.1.3__tar.gz

{types_array_api-1.1.1/src/types_array_api.egg-info → types_array_api-1.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: types-array-api
-Version: 1.1.1
+Version: 1.1.3
 Summary: Typing for array API and array-api-compat
 Author-email: 34j <34j.95a2p@simplelogin.com>
 License-Expression: MIT AND Apache-2.0
@@ -22,10 +22,11 @@ Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: LICENSE-MIT
-Requires-Dist: attrs>=25.3.0
-Requires-Dist: rich>=10
-Requires-Dist: typer<1,>=0.15
 Requires-Dist: typing-extensions>=4.14.0
+Provides-Extra: cli
+Requires-Dist: attrs>=25.3.0; extra == "cli"
+Requires-Dist: rich>=14; extra == "cli"
+Requires-Dist: typer>=0.16; extra == "cli"
 Dynamic: license-file
 
 # Python array API standard typing
@@ -82,7 +83,7 @@ pip install types-array-api
 
 ### Type stubs
 
-Provices type stubs for [`array-api-compat`](https://data-apis.org/array-api-compat/).
+Autocompletion for [`array-api-compat`](https://data-apis.org/array-api-compat/) is available in your IDE **just by installing** this package.
 
 ```python
 import array_api_compat
@@ -93,19 +94,31 @@ xp = array_api_compat.array_namespace(x)
 ![Screenshot 1](https://raw.githubusercontent.com/34j/array-api/main/docs/_static/screenshot1.png)
 ![Screenshot 2](https://raw.githubusercontent.com/34j/array-api/main/docs/_static/screenshot2.png)
 
-### Array Type
+### Typing functions using `Array`
 
-```python
-from array_api._2024_12 import Array
+There are multiple ways to type functions:
 
+- ```python
+  from array_api._2024_12 import Array
 
-def my_function[TArray: Array](x: TArray) -> TArray:
-    return x + 1
-```
+  def simple(x: Array) -> Array:
+      return x + 1
+  ```
+
+  The simplest way to enjoy autocompletion for `Array`. This should be enough for most use cases.
+
+- To make sure that the same type of array is returned (`ndarray`→`ndarray`, `Tensor`→`Tensor`), a `TypeVar` bound to `Array` can be used:
+
+  ```python
+  def generic[TArray: Array](x: TArray) -> TArray:
+      return x + 1
+  ```
+
+## Advanced Usage
 
 ### Namespace Type
 
-You can test if an object matches the Protocol by:
+You can test if an object matches the Protocol as they are [`runtime-checkable`](https://docs.python.org/3/library/typing.html#typing.runtime_checkable):
 
 ```python
 import array_api_strict
@@ -119,6 +132,47 @@ assert isinstance(array_api_strict, ArrayNamespace)
 assert not isinstance(array_api_strict, ArrayNamespaceFull)
 ```
 
+### Shape Typing
+
+- To clarify the input and output shapes, `ShapedArray` and `ShapedAnyArray` can be used:
+
+  ```python
+  from array_api._2024_12 import ShapedAnyArray as Array
+
+  def sum_last_axis[*TShape](x: Array[*TShape, Any]) -> Array[*TShape]:
+      return xp.sum(x, axis=-1)
+  ```
+
+  More complex example using [NewType](https://docs.python.org/3/library/typing.html#newtype) or [type aliases](https://docs.python.org/3/library/typing.html#type-aliases):
+
+  ```python
+  RTheta = NewType("RTheta", int)
+  XY = NewType("XY", int)
+  def polar_coordinates[*TShape](randtheta: Array[*TShape, RTheta]) -> Array[*TShape, XY]:
+      """Convert polar coordinates to Cartesian coordinates."""
+      r = randtheta[..., 0]
+      theta = randtheta[..., 1]
+      x = r * xp.cos(theta)
+      y = r * xp.sin(theta)
+      return xp.stack((x, y), axis=-1)
+  ```
+
+  Note that `ShapedAnyArray` exists only for **documentation purposes** and internally it is treated as `Array`.
+  Using both generic and shaped are impossible due to [python/typing#548](https://github.com/python/typing/issues/548).
+
+- Note that the below example is ideal but impossible due to Python specification.
+
+  ```python
+  def impossible[
+      TDtype,
+      TDevice,
+      *TShapeFormer: int,
+      *TShapeLatter: int,
+      TArray: Array
+  ](x: TArray[*TShapeFormer, *TShapeLatter | Literal[1], TDtype, TDevice], y: TArray[*TShapeLatter | Literal[1], TDtype, TDevice]) -> TArray[*TShapeFormer, *TShapeLatter, TDtype, TDevice]:
+      return x + y # broadcasting
+  ```
+
 ## Contributors ✨
 
 Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):

{types_array_api-1.1.1 → types_array_api-1.1.3}/README.md

@@ -52,7 +52,7 @@ pip install types-array-api
 
 ### Type stubs
 
-Provices type stubs for [`array-api-compat`](https://data-apis.org/array-api-compat/).
+Autocompletion for [`array-api-compat`](https://data-apis.org/array-api-compat/) is available in your IDE **just by installing** this package.
 
 ```python
 import array_api_compat
@@ -63,19 +63,31 @@ xp = array_api_compat.array_namespace(x)
 ![Screenshot 1](https://raw.githubusercontent.com/34j/array-api/main/docs/_static/screenshot1.png)
 ![Screenshot 2](https://raw.githubusercontent.com/34j/array-api/main/docs/_static/screenshot2.png)
 
-### Array Type
+### Typing functions using `Array`
 
-```python
-from array_api._2024_12 import Array
+There are multiple ways to type functions:
 
+- ```python
+  from array_api._2024_12 import Array
 
-def my_function[TArray: Array](x: TArray) -> TArray:
-    return x + 1
-```
+  def simple(x: Array) -> Array:
+      return x + 1
+  ```
+
+  The simplest way to enjoy autocompletion for `Array`. This should be enough for most use cases.
+
+- To make sure that the same type of array is returned (`ndarray`→`ndarray`, `Tensor`→`Tensor`), a `TypeVar` bound to `Array` can be used:
+
+  ```python
+  def generic[TArray: Array](x: TArray) -> TArray:
+      return x + 1
+  ```
+
+## Advanced Usage
 
 ### Namespace Type
 
-You can test if an object matches the Protocol by:
+You can test if an object matches the Protocol as they are [`runtime-checkable`](https://docs.python.org/3/library/typing.html#typing.runtime_checkable):
 
 ```python
 import array_api_strict
@@ -89,6 +101,47 @@ assert isinstance(array_api_strict, ArrayNamespace)
 assert not isinstance(array_api_strict, ArrayNamespaceFull)
 ```
 
+### Shape Typing
+
+- To clarify the input and output shapes, `ShapedArray` and `ShapedAnyArray` can be used:
+
+  ```python
+  from array_api._2024_12 import ShapedAnyArray as Array
+
+  def sum_last_axis[*TShape](x: Array[*TShape, Any]) -> Array[*TShape]:
+      return xp.sum(x, axis=-1)
+  ```
+
+  More complex example using [NewType](https://docs.python.org/3/library/typing.html#newtype) or [type aliases](https://docs.python.org/3/library/typing.html#type-aliases):
+
+  ```python
+  RTheta = NewType("RTheta", int)
+  XY = NewType("XY", int)
+  def polar_coordinates[*TShape](randtheta: Array[*TShape, RTheta]) -> Array[*TShape, XY]:
+      """Convert polar coordinates to Cartesian coordinates."""
+      r = randtheta[..., 0]
+      theta = randtheta[..., 1]
+      x = r * xp.cos(theta)
+      y = r * xp.sin(theta)
+      return xp.stack((x, y), axis=-1)
+  ```
+
+  Note that `ShapedAnyArray` exists only for **documentation purposes** and internally it is treated as `Array`.
+  Using both generic and shaped are impossible due to [python/typing#548](https://github.com/python/typing/issues/548).
+
+- Note that the below example is ideal but impossible due to Python specification.
+
+  ```python
+  def impossible[
+      TDtype,
+      TDevice,
+      *TShapeFormer: int,
+      *TShapeLatter: int,
+      TArray: Array
+  ](x: TArray[*TShapeFormer, *TShapeLatter | Literal[1], TDtype, TDevice], y: TArray[*TShapeLatter | Literal[1], TDtype, TDevice]) -> TArray[*TShapeFormer, *TShapeLatter, TDtype, TDevice]:
+      return x + y # broadcasting
+  ```
+
 ## Contributors ✨
 
 Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):

{types_array_api-1.1.1 → types_array_api-1.1.3}/pyproject.toml

@@ -4,7 +4,7 @@ requires = [ "setuptools" ]
 
 [project]
 name = "types-array-api"
-version = "1.1.1"
+version = "1.1.3"
 description = "Typing for array API and array-api-compat"
 readme = "README.md"
 license = "MIT AND Apache-2.0"
@@ -26,24 +26,30 @@ classifiers = [
 ]
 
 dependencies = [
-  "attrs>=25.3.0",
-  "rich>=10",
-  "typer>=0.15,<1",
   "typing-extensions>=4.14.0",
 ]
+optional-dependencies.cli = [
+  "attrs>=25.3.0",
+  "rich>=14",
+  "typer>=0.16",
+]
 urls."Bug Tracker" = "https://github.com/34j/types-array-api/issues"
 urls.Changelog = "https://github.com/34j/types-array-api/blob/main/CHANGELOG.md"
 urls.documentation = "https://array-api.readthedocs.io"
 urls.repository = "https://github.com/34j/types-array-api"
 scripts.array-api = "array_api.cli:app"
+scripts.types-array-api = "array_api.cli:app"
 
 [dependency-groups]
 dev = [
   "array-api-compat>=1.11.2",
   "array-api-strict>=2.3.1",
+  "attrs>=25.3.0",
   "pytest>=8,<9",
   "pytest-cov>=6,<7",
+  "rich>=14.0.0",
   "ruff>=0.11.13",
+  "typer>=0.16.0",
 ]
 docs = [
   "furo>=2023.5.20; python_version>='3.11'",

{types_array_api-1.1.1 → types_array_api-1.1.3}/src/array_api/_2022_12.py

@@ -5,9 +5,15 @@ from collections.abc import Buffer as SupportsBufferProtocol
 from collections.abc import Sequence
 from enum import Enum
 from types import EllipsisType as ellipsis
-from typing import Any, Literal, Protocol, runtime_checkable
+from typing import (
+    Any,
+    Literal,
+    Protocol,
+    runtime_checkable,
+)
 
 from typing_extensions import CapsuleType as PyCapsule
+from typing_extensions import Self
 
 inf = float("inf")
 
@@ -42,7 +48,7 @@ class finfo_object[TDtype](Protocol):
 
 
 @runtime_checkable
-class Array[TArray: Array, TDtype, TDevice](Protocol):
+class Array[TDtype, TDevice](Protocol):
     def __init__(self) -> None:
         """Initialize the attributes for the array object class."""
         ...
@@ -74,7 +80,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         ...
 
     @property
-    def mT(self) -> TArray:
+    def mT(self) -> Self:
         """
         Transpose of a matrix (or a stack of matrices).
 
@@ -142,7 +148,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         ...
 
     @property
-    def T(self) -> TArray:
+    def T(self) -> Self:
         """
         Transpose of the array.
 
@@ -160,7 +166,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __abs__(self, /) -> TArray:
+    def __abs__(self, /) -> Self:
         """
         Calculates the absolute value for each element of an array instance.
 
@@ -191,7 +197,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __add__(self, other: int | float | complex | TArray, /) -> TArray:
+    def __add__(self, other: int | float | complex | Self, /) -> Self:
         """
         Calculates the sum for each element of an array instance with the respective element of the array ``other``.
 
@@ -219,7 +225,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __and__(self, other: int | bool | TArray, /) -> TArray:
+    def __and__(self, other: int | bool | Self, /) -> Self:
         """
         Evaluates ``self_i & other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -242,7 +248,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __array_namespace__(self, /, *, api_version: str | None = None) -> Any:
+    def __array_namespace__(self, /, *, api_version: str | None = None) -> ArrayNamespace[Self, TDtype, TDevice]:
         """
         Returns an object that has all the array API functions on it.
 
@@ -427,7 +433,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __eq__(self, other: int | float | complex | bool | TArray, /) -> TArray:  # type: ignore[override]
+    def __eq__(self, other: int | float | complex | bool | Self, /) -> Self:  # type: ignore[override]
         """
         Computes the truth value of ``self_i == other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -485,7 +491,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __floordiv__(self, other: int | float | TArray, /) -> TArray:
+    def __floordiv__(self, other: int | float | Self, /) -> Self:
         """
         Evaluates ``self_i // other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -511,7 +517,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __ge__(self, other: int | float | TArray, /) -> TArray:
+    def __ge__(self, other: int | float | Self, /) -> Self:
         """
         Computes the truth value of ``self_i >= other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -537,7 +543,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __getitem__(self, key: int | slice | ellipsis | None | tuple[int | slice | ellipsis | None, ...] | TArray, /) -> TArray:
+    def __getitem__(self, key: int | slice | ellipsis | None | tuple[int | slice | ellipsis | None, ...] | Self, /) -> Self:
         """
         Returns ``self[key]``.
 
@@ -558,7 +564,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __gt__(self, other: int | float | TArray, /) -> TArray:
+    def __gt__(self, other: int | float | Self, /) -> Self:
         """
         Computes the truth value of ``self_i > other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -645,7 +651,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __invert__(self, /) -> TArray:
+    def __invert__(self, /) -> Self:
         """
         Evaluates ``~self_i`` for each element of an array instance.
 
@@ -666,7 +672,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __le__(self, other: int | float | TArray, /) -> TArray:
+    def __le__(self, other: int | float | Self, /) -> Self:
         """
         Computes the truth value of ``self_i <= other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -692,7 +698,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __lshift__(self, other: int | TArray, /) -> TArray:
+    def __lshift__(self, other: int | Self, /) -> Self:
         """
         Evaluates ``self_i << other_i`` for each element of an array instance with the respective element  of the array ``other``.
 
@@ -715,7 +721,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __lt__(self, other: int | float | TArray, /) -> TArray:
+    def __lt__(self, other: int | float | Self, /) -> Self:
         """
         Computes the truth value of ``self_i < other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -741,7 +747,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __matmul__(self, other: TArray, /) -> TArray:
+    def __matmul__(self, other: Self, /) -> Self:
         """
         Computes the matrix product.
 
@@ -791,7 +797,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __mod__(self, other: int | float | TArray, /) -> TArray:
+    def __mod__(self, other: int | float | Self, /) -> Self:
         """
         Evaluates ``self_i % other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -817,7 +823,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __mul__(self, other: int | float | complex | TArray, /) -> TArray:
+    def __mul__(self, other: int | float | complex | Self, /) -> Self:
         """
         Calculates the product for each element of an array instance with the respective element of the array ``other``.
 
@@ -848,7 +854,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __ne__(self, other: int | float | complex | bool | TArray, /) -> TArray:  # type: ignore[override]
+    def __ne__(self, other: int | float | complex | bool | Self, /) -> Self:  # type: ignore[override]
         """
         Computes the truth value of ``self_i != other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -876,7 +882,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __neg__(self, /) -> TArray:
+    def __neg__(self, /) -> Self:
         """
         Evaluates ``-self_i`` for each element of an array instance.
 
@@ -908,7 +914,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __or__(self, other: int | bool | TArray, /) -> TArray:
+    def __or__(self, other: int | bool | Self, /) -> Self:
         """
         Evaluates ``self_i | other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -933,7 +939,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __pos__(self, /) -> TArray:
+    def __pos__(self, /) -> Self:
         """
         Evaluates ``+self_i`` for each element of an array instance.
 
@@ -954,7 +960,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __pow__(self, other: int | float | complex | TArray, /) -> TArray:
+    def __pow__(self, other: int | float | complex | Self, /) -> Self:
         """
         Calculates an implementation-dependent approximation of exponentiation by raising each element (the base) of an array instance to the power of ``other_i`` (the exponent), where ``other_i`` is the corresponding element of the array ``other``.
 
@@ -987,7 +993,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __rshift__(self, other: int | TArray, /) -> TArray:
+    def __rshift__(self, other: int | Self, /) -> Self:
         """
         Evaluates ``self_i >> other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -1010,7 +1016,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __setitem__(self, key: int | slice | ellipsis | tuple[int | slice | ellipsis, ...] | TArray, value: int | float | complex | bool | TArray, /) -> None:
+    def __setitem__(self, key: int | slice | ellipsis | tuple[int | slice | ellipsis, ...] | Self, value: int | float | complex | bool | Self, /) -> None:
         """
         Sets ``self[key]`` to ``value``.
 
@@ -1037,7 +1043,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __sub__(self, other: int | float | complex | TArray, /) -> TArray:
+    def __sub__(self, other: int | float | complex | Self, /) -> Self:
         """
         Calculates the difference for each element of an array instance with the respective element of the array ``other``.
 
@@ -1067,7 +1073,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __truediv__(self, other: int | float | complex | TArray, /) -> TArray:
+    def __truediv__(self, other: int | float | complex | Self, /) -> Self:
         """
         Evaluates ``self_i / other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -1100,7 +1106,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def __xor__(self, other: int | bool | TArray, /) -> TArray:
+    def __xor__(self, other: int | bool | Self, /) -> Self:
         """
         Evaluates ``self_i ^ other_i`` for each element of an array instance with the respective element of the array ``other``.
 
@@ -1123,7 +1129,7 @@ class Array[TArray: Array, TDtype, TDevice](Protocol):
         """
         ...
 
-    def to_device(self, device: TDevice, /, *, stream: int | Any | None = None) -> TArray:
+    def to_device(self, device: TDevice, /, *, stream: int | Any | None = None) -> Self:
         """
         Copy the array from the device on which it currently resides to the specified ``device``.
 
@@ -7724,3 +7730,12 @@ class FftNamespace[TArray: Array, TDevice](Protocol):
 class ArrayNamespaceFull[TArray: Array, TDtype, TDevice](ArrayNamespace[TArray, TDtype, TDevice], Protocol):
     linalg: LinalgNamespace[TArray, TDtype]
     fft: FftNamespace[TArray, TDevice]
+
+
+@runtime_checkable
+class ShapedArray[*T, TDevice, TDtype](Array[TDevice, TDtype], Protocol):
+    @property
+    def shape(self) -> tuple[*T]: ...  # type: ignore[override]
+
+
+type ShapedAnyArray[*T] = ShapedArray[*T, Any, Any]