Parse and add type signatures to all the function stubs

This only adds parameter types. Return types will be added in the next commit.

Author: asmeurer

array_api_tests/function_stubs/_types.py

@@ -0,0 +1,15 @@
+"""
+This file defines the types for type annotations.
+
+The type variables should be replaced with the actual types for a given
+library, e.g., for NumPy TypeVar('array') would be replaced with ndarray.
+"""
+
+from typing import Literal, Optional, Tuple, Union, TypeVar
+
+array = TypeVar('array')
+device = TypeVar('device')
+dtype = TypeVar('dtype')
+
+__all__ = ['Literal', 'Optional', 'Tuple', 'Union', 'array', 'device', 'dtype']
+

array_api_tests/function_stubs/array_object.py

@@ -14,157 +14,159 @@
 2. There is no real way to test that anyway.
 """
 
-def __abs__(x):
+from ._types import array
+
+def __abs__(x: array):
     """
     Note: __abs__ is a method of the array object.
     """
     pass
 
-def __add__(x1, x2):
+def __add__(x1: array, x2: array):
     """
     Note: __add__ is a method of the array object.
     """
     pass
 
-def __and__(x1, x2):
+def __and__(x1: array, x2: array):
     """
     Note: __and__ is a method of the array object.
     """
     pass
 
-def __eq__(x1, x2):
+def __eq__(x1: array, x2: array):
     """
     Note: __eq__ is a method of the array object.
     """
     pass
 
-def __floordiv__(x1, x2):
+def __floordiv__(x1: array, x2: array):
     """
     Note: __floordiv__ is a method of the array object.
     """
     pass
 
-def __ge__(x1, x2):
+def __ge__(x1: array, x2: array):
     """
     Note: __ge__ is a method of the array object.
     """
     pass
 
-def __getitem__(x, key):
+def __ge__(x1: array, x2: array):
     """
     Note: __getitem__ is a method of the array object.
     """
     pass
 
-def __gt__(x1, x2):
+def __gt__(x1: array, x2: array):
     """
     Note: __gt__ is a method of the array object.
     """
     pass
 
-def __invert__(x):
+def __invert__(x: array):
     """
     Note: __invert__ is a method of the array object.
     """
     pass
 
-def __le__(x1, x2):
+def __le__(x1: array, x2: array):
     """
     Note: __le__ is a method of the array object.
     """
     pass
 
-def __len__(x):
+def __le__(x1: array, x2: array):
     """
     Note: __len__ is a method of the array object.
     """
     pass
 
-def __lshift__(x1, x2):
+def __lshift__(x1: array, x2: array):
     """
     Note: __lshift__ is a method of the array object.
     """
     pass
 
-def __lt__(x1, x2):
+def __lt__(x1: array, x2: array):
     """
     Note: __lt__ is a method of the array object.
     """
     pass
 
-def __matmul__(x1, x2):
+def __matmul__(x1: array, x2: array):
     """
     Note: __matmul__ is a method of the array object.
     """
     pass
 
-def __mod__(x1, x2):
+def __mod__(x1: array, x2: array):
     """
     Note: __mod__ is a method of the array object.
     """
     pass
 
-def __mul__(x1, x2):
+def __mul__(x1: array, x2: array):
     """
     Note: __mul__ is a method of the array object.
     """
     pass
 
-def __ne__(x1, x2):
+def __ne__(x1: array, x2: array):
     """
     Note: __ne__ is a method of the array object.
     """
     pass
 
-def __neg__(x):
+def __neg__(x: array):
     """
     Note: __neg__ is a method of the array object.
     """
     pass
 
-def __or__(x1, x2):
+def __or__(x1: array, x2: array):
     """
     Note: __or__ is a method of the array object.
     """
     pass
 
-def __pos__(x):
+def __pos__(x: array):
     """
     Note: __pos__ is a method of the array object.
     """
     pass
 
-def __pow__(x1, x2):
+def __pow__(x1: array, x2: array):
     """
     Note: __pow__ is a method of the array object.
     """
     pass
 
-def __rshift__(x1, x2):
+def __rshift__(x1: array, x2: array):
     """
     Note: __rshift__ is a method of the array object.
     """
     pass
 
-def __setitem__(x, key, value):
+def __rshift__(x1: array, x2: array):
     """
     Note: __setitem__ is a method of the array object.
     """
     pass
 
-def __sub__(x1, x2):
+def __sub__(x1: array, x2: array):
     """
     Note: __sub__ is a method of the array object.
     """
     pass
 
-def __truediv__(x1, x2):
+def __truediv__(x1: array, x2: array):
     """
     Note: __truediv__ is a method of the array object.
     """
     pass
 
-def __xor__(x1, x2):
+def __xor__(x1: array, x2: array):
     """
     Note: __xor__ is a method of the array object.
     """

array_api_tests/function_stubs/constants.py

@@ -14,6 +14,7 @@
 2. There is no real way to test that anyway.
 """
 
+
 e = None
 
 inf = None

array_api_tests/function_stubs/creation_functions.py

@@ -14,37 +14,39 @@
 2. There is no real way to test that anyway.
 """
 
-def arange(start, *, stop=None, step=1, dtype=None, device=None):
+from ._types import Optional, Tuple, Union, array, device, dtype
+
+def arange(start: Union[int, float], *, stop: Optional[Union[int, float]] = None, step: Union[int, float] = 1, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def empty(shape, *, dtype=None, device=None):
+def empty(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def empty_like(x, *, dtype=None, device=None):
+def empty_like(x: array, *, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def eye(N, *, M=None, k=0, dtype=None, device=None):
+def eye(N: int, *, M: Optional[int] = None, k: Optional[int] = 0, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def full(shape, fill_value, *, dtype=None, device=None):
+def full(shape: Union[int, Tuple[int, ...]], fill_value: Union[int, float], *, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def full_like(x, fill_value, *, dtype=None, device=None):
+def full_like(x: array, fill_value: Union[int, float], *, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def linspace(start, stop, num, *, dtype=None, device=None, endpoint=True):
+def linspace(start: Union[int, float], stop: Union[int, float], num: int, *, dtype: Optional[dtype] = None, device: Optional[device] = None, endpoint: Optional[bool] = True):
     pass
 
-def ones(shape, *, dtype=None, device=None):
+def ones(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def ones_like(x, *, dtype=None, device=None):
+def ones_like(x: array, *, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def zeros(shape, *, dtype=None, device=None):
+def zeros(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
-def zeros_like(x, *, dtype=None, device=None):
+def zeros_like(x: array, *, dtype: Optional[dtype] = None, device: Optional[device] = None):
     pass
 
 __all__ = ['arange', 'empty', 'empty_like', 'eye', 'full', 'full_like', 'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like']