Add isfinite, isinf, isscalar, diagonal, allclose, and improve isnan support in NumPy stubs

also add support for scalar + array

fixes #209

Author: tadeu

numpy-stubs/__init__.pyi

@@ -390,6 +390,8 @@ class ndarray(Generic[_DType]):
     def __radd__(self, value: ndarray[_DType]) -> ndarray[_DType]: ...
     @overload
     def __radd__(self, value: _DType) -> ndarray[_DType]: ...
+    @overload
+    def __radd__(self, value: float) -> ndarray[_DType]: ...
     def __rand__(self, value: object) -> ndarray[_DType]: ...
     def __rdivmod__(self, value: object) -> Tuple[ndarray[_DType], ndarray[_DType]]: ...
     def __rfloordiv__(self, value: object) -> ndarray[_DType]: ...
@@ -774,10 +776,6 @@ def interp(
 ) -> ndarray: ...
 def isin(element: Sequence[_DType], test_element: _DType) -> ndarray[_DType]: ...
 @overload
-def isnan(x: float64) -> bool: ...
-@overload
-def isnan(x: ndarray[_DType]) -> ndarray[bool_]: ...
-@overload
 def ix_(x: ndarray[_DType]) -> ndarray[_DType]: ...
 @overload
 def ix_(x1: ndarray[_DType], x2: ndarray[_DType]) -> Tuple[ndarray[_DType], ndarray[_DType]]: ...
@@ -993,6 +991,92 @@ def set_printoptions(
     *,
     legacy: Any = ...,
 ) -> None: ...
+def isscalar(element: Any) -> bool: ...
+def diagonal(a: _ArrayLike, offset: int = ..., axis1: int = ..., axis2: int = ...) -> ndarray: ...
+def allclose(
+    a: Union[_ArrayLike, _FloatLike],
+    b: Union[_ArrayLike, _FloatLike],
+    rtol: float = ...,
+    atol: float = ...,
+    equal_nan: bool = ...,
+) -> bool: ...
+
+#
+# ufunc
+#
+
+# Backported from latest NumPy
+class ufunc:
+    @property
+    def __name__(self) -> str: ...
+    def __call__(
+        self,
+        *args: Union[_FloatLike, _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 = ...,
+        casting: Any = ...,
+        order: Any = ...,
+        dtype: Any = ...,
+        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]] = ...,
+    ) -> Any: ...
+    @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: ...
+
+isfinite: ufunc
+isinf: ufunc
+isnan: ufunc
 
 #
 # Specific values

tests/numpy_test.py

@@ -259,3 +259,52 @@ def test_interp() -> None:
 def test_genfromtxt() -> None:
     result = np.genfromtxt(["0.1, 0.2"], dtype=np.float64, delimiter=",")
     assert list(result) == [0.1, 0.2]
+
+
+def test_isfinite_isinf_isnan() -> None:
+    import math
+
+    assert np.isfinite(0.0)
+    assert not np.isfinite(np.inf)
+    assert np.isinf(np.inf)
+    assert not np.isfinite(math.inf)
+    assert not np.isfinite(np.nan)
+    assert np.isnan(np.nan)
+    assert np.all(np.isfinite([0.0, -np.inf]) == [True, False])
+    assert np.all(np.isfinite(np.array([0.0, np.nan])) == np.array([True, False]))
+    assert np.all(
+        np.isfinite(np.array([np.inf, np.nan], dtype=np.float32)) == np.array([False, False])
+    )
+    assert np.all(np.isnan([0.0, -np.inf]) == [False, False])
+    assert np.all(np.isinf([0.0, -np.inf]) == [False, True])
+
+
+def test_diagonal() -> None:
+    assert np.all(np.diagonal([[1]]) == np.array([1]))
+    x = np.arange(12).reshape(3, 4)
+    assert np.all(np.diagonal(x) == np.array([0.0, 5.0, 10.0]))
+
+
+def test_allclose() -> None:
+    assert np.allclose([1.0, 2.0], [1.0 + 1e-9, 2.0 + 1e-9])
+    assert np.allclose(np.array([1.0, 2.0]), np.array([1.0 + 1e-9, 2.0 + 1e-9]))
+    assert np.allclose(np.array([1.0, 1.0]), 1.0 + 1e-9)
+    assert np.allclose(1.0 + 1e-9, np.array([1.0, 1.0]))
+
+
+def test_isscalar() -> None:
+    assert np.isscalar(1.0)
+    assert not np.isscalar([1.0])
+    assert not np.isscalar(np.array([1.0]))
+    assert not np.isscalar(np.array([]))
+    assert np.isscalar(np.array([1.0], dtype=np.float32)[0])
+
+
+def test_newaxis() -> None:
+    x = np.array([1.0, 2.0])
+    assert x[np.newaxis, :].shape == (1, 2)
+
+
+def test_sum_scalar_before() -> None:
+    x = 273.15 + np.array([-0.1e2, -0.77e1])
+    assert isinstance(x, np.ndarray)