Add type annotations to the array api submodule function definitions

Some stubs still need to be modified to properly pass mypy type checking.
Also, 'device' is just left as a TypeVar() for now.

Author: asmeurer

numpy/_array_api/_creation_functions.py

@@ -1,6 +1,10 @@
+from __future__ import annotations
+
+from ._types import Optional, Tuple, Union, array, device, dtype
+
 import numpy as np
 
-def arange(start, /, *, stop=None, step=1, dtype=None, device=None):
+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) -> array:
     """
     Array API compatible wrapper for :py:func:`np.arange <numpy.arange>`.
 
@@ -11,7 +15,7 @@ def arange(start, /, *, stop=None, step=1, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.arange(start, stop=stop, step=step, dtype=dtype)
 
-def empty(shape, /, *, dtype=None, device=None):
+def empty(shape: Union[int, Tuple[int, ...]], /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.empty <numpy.empty>`.
 
@@ -22,7 +26,7 @@ def empty(shape, /, *, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.empty(shape, dtype=dtype)
 
-def empty_like(x, /, *, dtype=None, device=None):
+def empty_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.empty_like <numpy.empty_like>`.
 
@@ -33,7 +37,7 @@ def empty_like(x, /, *, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.empty_like(x, dtype=dtype)
 
-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) -> array:
     """
     Array API compatible wrapper for :py:func:`np.eye <numpy.eye>`.
 
@@ -44,7 +48,7 @@ def eye(N, /, *, M=None, k=0, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.eye(N, M=M, k=k, dtype=dtype)
 
-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) -> array:
     """
     Array API compatible wrapper for :py:func:`np.full <numpy.full>`.
 
@@ -55,7 +59,7 @@ def full(shape, fill_value, /, *, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.full(shape, fill_value, dtype=dtype)
 
-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) -> array:
     """
     Array API compatible wrapper for :py:func:`np.full_like <numpy.full_like>`.
 
@@ -66,7 +70,7 @@ def full_like(x, fill_value, /, *, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.full_like(x, fill_value, dtype=dtype)
 
-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) -> array:
     """
     Array API compatible wrapper for :py:func:`np.linspace <numpy.linspace>`.
 
@@ -77,7 +81,7 @@ def linspace(start, stop, num, /, *, dtype=None, device=None, endpoint=True):
         raise NotImplementedError("Device support is not yet implemented")
     return np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint)
 
-def ones(shape, /, *, dtype=None, device=None):
+def ones(shape: Union[int, Tuple[int, ...]], /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.ones <numpy.ones>`.
 
@@ -88,7 +92,7 @@ def ones(shape, /, *, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.ones(shape, dtype=dtype)
 
-def ones_like(x, /, *, dtype=None, device=None):
+def ones_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.ones_like <numpy.ones_like>`.
 
@@ -99,7 +103,7 @@ def ones_like(x, /, *, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.ones_like(x, dtype=dtype)
 
-def zeros(shape, /, *, dtype=None, device=None):
+def zeros(shape: Union[int, Tuple[int, ...]], /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.zeros <numpy.zeros>`.
 
@@ -110,7 +114,7 @@ def zeros(shape, /, *, dtype=None, device=None):
         raise NotImplementedError("Device support is not yet implemented")
     return np.zeros(shape, dtype=dtype)
 
-def zeros_like(x, /, *, dtype=None, device=None):
+def zeros_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.zeros_like <numpy.zeros_like>`.