From 3fb733e3b368f990b0119a7cad86aba2cffc5345 Mon Sep 17 00:00:00 2001 From: "Edson Tadeu M. Manoel" Date: Tue, 22 Dec 2020 11:37:12 -0300 Subject: [PATCH 1/2] Add `isfinite`, `isinf`, `isscalar`, `diagonal`, `allclose`, and improve `isnan` support in NumPy stubs also add support for scalar + array fixes #209 --- numpy-stubs/__init__.pyi | 92 ++++++++++++++++++++++++++++++++++++++-- tests/numpy_test.py | 50 ++++++++++++++++++++++ 2 files changed, 138 insertions(+), 4 deletions(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index c7c9909..de108fd 100644 --- a/numpy-stubs/__init__.pyi +++ b/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 diff --git a/tests/numpy_test.py b/tests/numpy_test.py index 663a6b5..cc67f77 100644 --- a/tests/numpy_test.py +++ b/tests/numpy_test.py @@ -259,3 +259,53 @@ 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) + assert x.dtype == np.float64 From e4f26eebd9d6d8dc4fcdede796b4b2712b1e4b15 Mon Sep 17 00:00:00 2001 From: "Edson Tadeu M. Manoel" Date: Tue, 22 Dec 2020 15:21:12 -0300 Subject: [PATCH 2/2] Improve scalar + array, and add complex support --- numpy-stubs/__init__.pyi | 19 +++++++++++++++---- tests/numpy_test.py | 10 ++++++++-- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index de108fd..a113b43 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -42,12 +42,17 @@ class void: # and this is why we let int32 be a subclass of int64; and similarly for float32 and float64 # the same logic applies when adding unsigned and signed values (uint + int -> int) +class complex128(void, int): + def __complex__(self) -> complex: ... + +class complex64(complex128): ... + # this would be the correct definition, but it makes `int` conflict with `float` # class float64(void, float): ... -class float64(void, int): +class float64(complex128): def __float__(self) -> float: ... -class float32(float64): ... +class float32(float64, complex64): ... class float16(float32): ... floating = float64 @@ -70,6 +75,8 @@ integer = int64 _DType = TypeVar( "_DType", bool_, + complex64, + complex128, float16, float32, float64, @@ -87,6 +94,8 @@ _DType = TypeVar( _DType2 = TypeVar( "_DType2", bool_, + complex64, + complex128, float16, float32, float64, @@ -110,7 +119,7 @@ _ScalarLike = Union[_DType, str, int, float] _ConditionType = Union[ndarray[bool_], bool_, bool] newaxis: None = ... -_AnyNum = Union[int, float, bool] +_AnyNum = Union[int, float, bool, complex] # generic types that are only allowed to take on dtype values _Float = TypeVar("_Float", float16, float32, float64) @@ -391,7 +400,7 @@ class ndarray(Generic[_DType]): @overload def __radd__(self, value: _DType) -> ndarray[_DType]: ... @overload - def __radd__(self, value: float) -> ndarray[_DType]: ... + def __radd__(self, value: float) -> ndarray[_DType2]: ... def __rand__(self, value: object) -> ndarray[_DType]: ... def __rdivmod__(self, value: object) -> Tuple[ndarray[_DType], ndarray[_DType]]: ... def __rfloordiv__(self, value: object) -> ndarray[_DType]: ... @@ -506,6 +515,8 @@ def array(object: _NestedList[int]) -> ndarray[int64]: ... @overload def array(object: _NestedList[float]) -> ndarray[float64]: ... @overload +def array(object: _NestedList[complex]) -> ndarray[complex64]: ... +@overload def array(object: _NestedList[str]) -> ndarray[str_]: ... @overload def array(object: str) -> ndarray[str_]: ... diff --git a/tests/numpy_test.py b/tests/numpy_test.py index cc67f77..82f6cf4 100644 --- a/tests/numpy_test.py +++ b/tests/numpy_test.py @@ -8,6 +8,8 @@ DType = TypeVar( "DType", np.bool_, + np.complex64, + np.complex128, np.float32, np.float64, np.int8, @@ -306,6 +308,10 @@ def test_newaxis() -> None: def test_sum_scalar_before() -> None: - x = 273.15 + np.array([-0.1e2, -0.77e1]) + x: np.ndarray[np.float64] = 273.15 + np.array([10, 20]) assert isinstance(x, np.ndarray) - assert x.dtype == np.float64 + assert_dtype(x, np.float64) + + y: np.ndarray[np.complex128] = 10.0 + np.array([1j, 2j]) + assert isinstance(y, np.ndarray) + assert_dtype(y, np.complex128)