ENH: add basic typing for np.ufunc (#44)

* ENH: add basic typing for `np.ufunc`

This adds basic type hints for `np.ufunc` and types all ufuncs in the
top-level namespace as `np.ufunc`. Ufuncs are highly dynamic, e.g.

- Their call signatures can vary
- The `reduce`/`accumulate`/`reduceat`/`outer`/`at` methods may or may
  not always raise

so it is difficult to have precise types. A path forward would be
writing a mypy plugin to get more precise typing.

* MAINT: fix return type for `ufunc.__call__`

Co-authored-by: Stephan Hoyer <shoyer@gmail.com>

.gitignore

@@ -4,4 +4,5 @@ __pycache__
 numpy_stubs.egg-info/
 venv
 .idea
+*~
 **~

numpy-stubs/__init__.pyi

@@ -6,6 +6,7 @@ from numpy.core._internal import _ctypes
 from typing import (
     Any,
     ByteString,
+    Callable,
     Container,
     Callable,
     Dict,
@@ -618,5 +619,160 @@ WRAP: int
 little_endian: int
 tracemalloc_domain: int
 
+class ufunc:
+    def __call__(
+        self,
+        *args: _ArrayLike,
+        out: Optional[Union[ndarray, Tuple[ndarray, ...]]] = ...,
+        where: Optional[ndarray] = ...,
+        # The list should be a list of tuples of ints, but since we
+        # don't know the signature it would need to be
+        # Tuple[int, ...]. But, since List is invariant something like
+        # e.g. List[Tuple[int, int]] isn't a subtype of
+        # List[Tuple[int, ...]], so we can't type precisely here.
+        axes: List[Any] = ...,
+        axis: int = ...,
+        keepdims: bool = ...,
+        # TODO: make this precise when we can use Literal.
+        casting: str = ...,
+        # TODO: make this precise when we can use Literal.
+        order: Optional[str] = ...,
+        dtype: Optional[_DtypeLike] = ...,
+        subok: bool = ...,
+        signature: Union[str, Tuple[str]] = ...,
+        # In reality this should be a length of list 3 containing an
+        # int, an int, and a callable, but there's no way to express
+        # that.
+        extobj: List[Union[int, Callable]] = ...,
+    ) -> Union[ndarray, generic]: ...
+    @property
+    def nin(self) -> int: ...
+    @property
+    def nout(self) -> int: ...
+    @property
+    def nargs(self) -> int: ...
+    @property
+    def ntypes(self) -> int: ...
+    @property
+    def types(self) -> List[str]: ...
+    # Broad return type because it has to encompass things like
+    #
+    # >>> np.logical_and.identity is True
+    # True
+    # >>> np.add.identity is 0
+    # True
+    # >>> np.sin.identity is None
+    # True
+    #
+    # and any user-defined ufuncs.
+    @property
+    def identity(self) -> Any: ...
+    # This is None for ufuncs and a string for gufuncs.
+    @property
+    def signature(self) -> Optional[str]: ...
+    # The next four methods will always exist, but they will just
+    # raise a ValueError ufuncs with that don't accept two input
+    # arguments and return one output argument. Because of that we
+    # can't type them very precisely.
+    @property
+    def reduce(self) -> Any: ...
+    @property
+    def accumulate(self) -> Any: ...
+    @property
+    def reduceat(self) -> Any: ...
+    @property
+    def outer(self) -> Any: ...
+    # Similarly at won't be defined for ufuncs that return multiple
+    # outputs, so we can't type it very precisely.
+    @property
+    def at(self) -> Any: ...
+
+absolute: ufunc
+add: ufunc
+arccos: ufunc
+arccosh: ufunc
+arcsin: ufunc
+arcsinh: ufunc
+arctan2: ufunc
+arctan: ufunc
+arctanh: ufunc
+bitwise_and: ufunc
+bitwise_or: ufunc
+bitwise_xor: ufunc
+cbrt: ufunc
+ceil: ufunc
+conjugate: ufunc
+copysign: ufunc
+cos: ufunc
+cosh: ufunc
+deg2rad: ufunc
+degrees: ufunc
+divmod: ufunc
+equal: ufunc
+exp2: ufunc
+exp: ufunc
+expm1: ufunc
+fabs: ufunc
+float_power: ufunc
+floor: ufunc
+floor_divide: ufunc
+fmax: ufunc
+fmin: ufunc
+fmod: ufunc
+frexp: ufunc
+gcd: ufunc
+greater: ufunc
+greater_equal: ufunc
+heaviside: ufunc
+hypot: ufunc
+invert: ufunc
+isfinite: ufunc
+isinf: ufunc
+isnan: ufunc
+isnat: ufunc
+lcm: ufunc
+ldexp: ufunc
+left_shift: ufunc
+less: ufunc
+less_equal: ufunc
+log10: ufunc
+log1p: ufunc
+log2: ufunc
+log: ufunc
+logaddexp2: ufunc
+logaddexp: ufunc
+logical_and: ufunc
+logical_not: ufunc
+logical_or: ufunc
+logical_xor: ufunc
+matmul: ufunc
+maximum: ufunc
+minimum: ufunc
+modf: ufunc
+multiply: ufunc
+negative: ufunc
+nextafter: ufunc
+not_equal: ufunc
+positive: ufunc
+power: ufunc
+rad2deg: ufunc
+radians: ufunc
+reciprocal: ufunc
+remainder: ufunc
+right_shift: ufunc
+rint: ufunc
+sign: ufunc
+signbit: ufunc
+sin: ufunc
+sinh: ufunc
+spacing: ufunc
+sqrt: ufunc
+square: ufunc
+subtract: ufunc
+tan: ufunc
+tanh: ufunc
+true_divide: ufunc
+trunc: ufunc
+
 # TODO(shoyer): remove when the full numpy namespace is defined
 def __getattr__(name: str) -> Any: ...

scripts/find_ufuncs.py

@@ -0,0 +1,21 @@
+import numpy as np
+
+
+def main():
+    ufuncs = []
+    for obj_name in np.__dir__():
+        obj = getattr(np, obj_name)
+        if isinstance(obj, np.ufunc):
+            ufuncs.append(obj)
+
+    ufunc_stubs = []
+    for ufunc in set(ufuncs):
+        ufunc_stubs.append(f'{ufunc.__name__}: ufunc')
+    ufunc_stubs.sort()
+
+    for stub in ufunc_stubs:
+        print(stub)
+
+
+if __name__ == '__main__':
+    main()

tests/fail/ufuncs.py

@@ -0,0 +1,5 @@
+import numpy as np
+
+np.sin.nin + 'foo'  # E: Unsupported operand types
+np.sin(1, foo='bar')  # E: Unexpected keyword argument
+np.sin(1, extobj=['foo', 'foo', 'foo'])  # E: incompatible type

tests/pass/ufuncs.py

@@ -0,0 +1,14 @@
+import numpy as np
+
+np.sin(1)
+np.sin([1, 2, 3])
+np.sin(1, out=np.empty(1))
+np.matmul(
+    np.ones((2, 2, 2)),
+    np.ones((2, 2, 2)),
+    axes=[(0, 1), (0, 1), (0, 1)],
+)
+np.sin(1, signature='D')
+np.sin(1, extobj=[16, 1, lambda: None])
+np.sin(1) + np.sin(1)
+np.sin.types[0]