Fill gaps in numpy scalar operations and function arguments

tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi

@@ -6,7 +6,17 @@
 from builtins import bool as py_bool
 from collections.abc import Iterator, Sequence
 from types import EllipsisType
-from typing import Any, Final, Generic, Literal, Never, overload, SupportsIndex, TypeVar
+from typing import (
+    Any,
+    Final,
+    Generic,
+    Literal,
+    Never,
+    overload,
+    Self,
+    SupportsIndex,
+    TypeVar,
+)
 
 import shape_extensions
 from numpy.__config__ import (
@@ -206,7 +216,7 @@ from numpy._shapes import (
     vecdot_shape,
     vecmat_shape,
 )
-from numpy._typing import _ArrayLikeInt_co, ArrayLike
+from numpy._typing import _ArrayLikeInt_co, ArrayLike, DTypeLike
 from numpy._typing._extended_precision import (
     complex256 as complex256,
     float128 as float128,
@@ -429,25 +439,75 @@ _ShapeT_co = TypeVar("_ShapeT_co", bound=_Shape, default=_Shape, covariant=True)
 type _Axis = int | tuple[int, ...] | None
 type _SingleAxis = int | None
 type _BasicIndex = int | slice | list[int] | None | EllipsisType
+type _NumberScalar = number | bool_ | int | float | complex
 
 class generic:
+    @property
+    def shape(self) -> tuple[()]: ...
+    def __bool__(self, /) -> py_bool: ...
+    @overload
+    def __array__(self, dtype: None = None, /) -> ndarray[[], _dtype[Self]]: ...
+    @overload
+    def __array__[DType: _dtype](self, dtype: DType, /) -> ndarray[[], DType]: ...
     def astype(self, dtype: Any, **kwargs: Any) -> Any: ...
 
-class number(generic): ...
-class integer(number): ...
+class number(generic):
+    def __neg__(self) -> Self: ...
+    def __pos__(self) -> Self: ...
+    def __abs__(self) -> number: ...
+    # The result dtype depends on both operands.
+    def __add__(self, other: _NumberScalar, /) -> number: ...
+    def __radd__(self, other: _NumberScalar, /) -> number | complex: ...
+    def __sub__(self, other: _NumberScalar, /) -> number: ...
+    def __rsub__(self, other: _NumberScalar, /) -> number | complex: ...
+    def __mul__(self, other: _NumberScalar, /) -> number: ...
+    def __rmul__(self, other: _NumberScalar, /) -> number | complex: ...
+    def __truediv__(self, other: _NumberScalar, /) -> number: ...
+    def __rtruediv__(self, other: _NumberScalar, /) -> number | complex: ...
+    def __pow__(self, other: _NumberScalar, mod: None = None, /) -> number: ...
+    def __rpow__(
+        self, other: _NumberScalar, mod: None = None, /
+    ) -> number | complex: ...
+    def __lt__(self, other: _NumberScalar, /) -> bool_: ...
+    def __le__(self, other: _NumberScalar, /) -> bool_: ...
+    def __gt__(self, other: _NumberScalar, /) -> bool_: ...
+    def __ge__(self, other: _NumberScalar, /) -> bool_: ...
+    def __int__(self, /) -> int: ...
+    def __float__(self, /) -> float: ...
 
-class signedinteger(integer):
+class integer(number):
+    def __index__(self, /) -> int: ...
     # TODO(stroxler): Sequence inputs to scalar constructors can produce arrays.
     def __init__(self, value: object = ...) -> None: ...
 
+class signedinteger(integer): ...
 class unsignedinteger(integer): ...
 class inexact(number): ...
 
 class floating(inexact):
+    def __init__(
+        self, value: number | bool_ | int | float | str | bytes = ...
+    ) -> None: ...
+
+class complexfloating(inexact):
+    def __init__(self, value: _NumberScalar | str | bytes = ...) -> None: ...
+    def __complex__(self, /) -> complex: ...
+
+class bool_(generic):
+    def __init__(self, value: _NumberScalar | str | bytes = ...) -> None: ...
+    def __int__(self, /) -> int: ...
     def __float__(self, /) -> float: ...
+    def __add__(self, other: _NumberScalar, /) -> number | bool_: ...
+    def __radd__(self, other: _NumberScalar, /) -> number | bool_: ...
+    def __sub__(self, other: number | int | float | complex, /) -> number: ...
+    def __rsub__(self, other: number | int | float | complex, /) -> number: ...
+    def __mul__(self, other: _NumberScalar, /) -> number | bool_: ...
+    def __rmul__(self, other: _NumberScalar, /) -> number | bool_: ...
+    def __truediv__(self, other: _NumberScalar, /) -> number: ...
+    def __rtruediv__(self, other: _NumberScalar, /) -> number: ...
+    def __pow__(self, other: _NumberScalar, mod: None = None, /) -> number: ...
+    def __rpow__(self, other: _NumberScalar, mod: None = None, /) -> number: ...
 
-class complexfloating(inexact): ...
-class bool_(generic): ...
 class float16(floating): ...
 class float32(floating): ...
 class float64(floating): ...
@@ -592,6 +652,7 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]):
     @overload
     def __iter__[N: IntVar](self: ndarray[[N], _DTypeT_co], /) -> Iterator[Any]: ...
     def __int__(self: ndarray[[]], /) -> int: ...
+    def __index__(self: ndarray[[], _dtype[integer]], /) -> int: ...
     def __float__(self: ndarray[[]], /) -> float: ...
     def __complex__(self: ndarray[[]], /) -> complex: ...
     # Only 2-D transpose is modeled for the NumPy shape-stub MVP.
@@ -600,30 +661,54 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]):
         self: ndarray[[N, P], _DTypeT_co],
     ) -> ndarray[[P, N], _DTypeT_co]: ...
     # Unpacking preserves gradual array shapes instead of using the scalar default.
+    @overload
     def __add__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __add__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __radd__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __radd__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __sub__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __sub__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __rsub__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __rsub__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __truediv__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __truediv__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __rtruediv__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __rtruediv__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __mul__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __mul__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __rmul__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __rmul__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
     def __floordiv__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
@@ -639,12 +724,18 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]):
     def __neg__(self) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
     def __pos__(self) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
     def __abs__(self, /) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __pow__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __pow__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
+    @overload
     def __rpow__[OtherShape: _Shape = []](
         self, other: ndarray[[*OtherShape]] | int | float
     ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ...
+    @overload
+    def __rpow__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ...
     # Scalar results also satisfy Python's truth-valued comparison protocols.
     @overload
     def __lt__(self: ndarray[[]], other: Any, /) -> Any: ...
@@ -745,12 +836,14 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]):
         axis: Axis = None,
         *,
         keepdims: KeepDims = False,
+        initial: _ArrayScalar = ...,
     ) -> ndarray[reduce_shape(_ShapeT_co, Axis, KeepDims), _DTypeT_co]: ...
     def max[Axis: Flag[_Axis], KeepDims: Flag[py_bool]](
         self,
         axis: Axis = None,
         *,
         keepdims: KeepDims = False,
+        initial: _ArrayScalar = ...,
     ) -> ndarray[reduce_shape(_ShapeT_co, Axis, KeepDims), _DTypeT_co]: ...
     def all[Axis: Flag[_Axis], KeepDims: Flag[py_bool]](
         self,
@@ -1626,7 +1719,7 @@ def asarray(
     copy: py_bool | None = ...,
     like: Any,
 ) -> Any: ...
-def round[Shape: _Shape](x: ndarray[Shape]) -> ndarray[Shape]: ...
+def round[Shape: _Shape](a: ndarray[Shape], decimals: int = 0) -> ndarray[Shape]: ...
 
 # Interpolation tracks query shapes and leaves the result dtype gradual.
 @overload
@@ -1676,7 +1769,25 @@ def diag[M: IntVar, N: IntVar, DType, K: Flag[int] = 0](
 def diag[S: _Shape, DType](
     v: ndarray[S, DType], k: int = 0
 ) -> ndarray[IntTuple, DType]: ...
-def arange[N: IntVar](stop: Int[N], /) -> ndarray[[N], dtype[intp]]: ...
+@overload
+def arange[N: IntVar](
+    stop: Int[N],
+    /,
+    *,
+    dtype: None = None,
+    device: Literal["cpu"] | None = None,
+    like: None = None,
+) -> ndarray[[N], dtype[intp]]: ...
+@overload
+def arange(
+    start: float,
+    stop: float | None = None,
+    step: float = 1,
+    dtype: DTypeLike | None = None,
+    *,
+    device: Literal["cpu"] | None = None,
+    like: None = None,
+) -> ndarray[[int]]: ...
 def expand_dims[Shape: _Shape, DType, Axis: Flag[int]](
     a: ndarray[Shape, DType], axis: Axis
 ) -> ndarray[expand_dims_shape(Shape, Axis), DType]: ...
@@ -1705,11 +1816,19 @@ def min[
     Axis: Flag[_Axis],
     KeepDims: Flag[py_bool],
 ](
-    a: ndarray[Shape, DType], axis: Axis = None, *, keepdims: KeepDims = False
+    a: ndarray[Shape, DType],
+    axis: Axis = None,
+    *,
+    keepdims: KeepDims = False,
+    initial: _ArrayScalar = ...,
 ) -> ndarray[reduce_shape(Shape, Axis, KeepDims), Any]: ...
 @overload
 def min(
-    a: Sequence[ArrayLike], axis: _Axis = None, *, keepdims: py_bool = False
+    a: Sequence[ArrayLike],
+    axis: _Axis = None,
+    *,
+    keepdims: py_bool = False,
+    initial: _ArrayScalar = ...,
 ) -> Any: ...
 @overload
 def max[
@@ -1718,11 +1837,19 @@ def max[
     Axis: Flag[_Axis],
     KeepDims: Flag[py_bool],
 ](
-    a: ndarray[Shape, DType], axis: Axis = None, *, keepdims: KeepDims = False
+    a: ndarray[Shape, DType],
+    axis: Axis = None,
+    *,
+    keepdims: KeepDims = False,
+    initial: _ArrayScalar = ...,
 ) -> ndarray[reduce_shape(Shape, Axis, KeepDims), Any]: ...
 @overload
 def max(
-    a: Sequence[ArrayLike], axis: _Axis = None, *, keepdims: py_bool = False
+    a: Sequence[ArrayLike],
+    axis: _Axis = None,
+    *,
+    keepdims: py_bool = False,
+    initial: _ArrayScalar = ...,
 ) -> Any: ...
 @overload
 def argmin[

tensor-shapes/pyrefly-numpy-stubs/test/test_arithmetic.py

@@ -53,6 +53,31 @@ def test_unary_arithmetic() -> None:
     assert_shape((-np.ones((3, 4))).shape, (3, 4))
 
 
+def test_scalar_array_arithmetic() -> None:
+    a = np.full((2, 3), 2.0, dtype=np.float32)
+
+    assert_shape((a + 1j).shape, (2, 3))
+    assert_shape((1j + a).shape, (2, 3))
+    assert_shape((a - 1j).shape, (2, 3))
+    assert_shape((1j - a).shape, (2, 3))
+    assert_shape((a * 1j).shape, (2, 3))
+    assert_shape((1j * a).shape, (2, 3))
+    assert_shape((a / 1j).shape, (2, 3))
+    assert_shape((1j / a).shape, (2, 3))
+    assert_shape((a**1j).shape, (2, 3))
+    assert_shape((1j**a).shape, (2, 3))
+    assert_type(a + 1j, np.ndarray[[2, 3]])
+    assert (a + 1j).dtype == np.complex64
+
+    scalar = np.array(2.0)
+    assert_shape((scalar * 1j).shape, ())
+
+    real = np.float64(2.0)
+    assert_shape((real * a).shape, (2, 3))
+    assert_shape((a * real).shape, (2, 3))
+    assert (a * real).dtype == np.float64
+
+
 def test_builtin_abs_preserves_shape() -> None:
     scalar = np.array(-2.0)
     vector = np.full(5, -1.0)

tensor-shapes/pyrefly-numpy-stubs/test/test_dtype_properties.py

@@ -8,7 +8,7 @@
 from typing import Any, assert_type
 
 import numpy as np
-from shape_extensions import assert_shape
+from shape_extensions import assert_raises, assert_shape
 
 
 def check_number_is_not_gradual() -> None:
@@ -64,6 +64,62 @@ def test_floating_and_complex_scalar_hierarchy() -> None:
     assert np.cdouble is np.complex128
 
 
+def test_numeric_scalar_operations() -> None:
+    integer = np.int64(2)
+    real = np.float64(2.0)
+    complex_ = np.complex128(1.0 + 2.0j)
+
+    assert_type(int(integer), int)
+    assert_type(integer.__index__(), int)
+    assert_type(float(real), float)
+    assert_type(complex(complex_), complex)
+    assert_shape((real**2).shape, ())
+    assert_shape((real * 1j).shape, ())
+    assert_type(1j * real, np.number | complex)
+    assert_shape((complex_ / integer).shape, ())
+    assert_shape(abs(complex_).shape, ())
+    assert int(np.uint64(2)) == 2
+    assert [10, 20, 30][integer] == 30
+    assert_type(max([real, real]), np.float64)
+    assert max([real, real + 1]) == 3.0
+    scalar_array = real.__array__()
+    assert_shape(scalar_array.shape, ())
+    assert_type(scalar_array.dtype, np.dtype[np.float64])
+    converted = real.__array__(np.dtype(np.complex128))
+    assert_shape(converted.shape, ())
+    assert_type(converted.dtype, np.dtype[np.complex128])
+
+    boolean = np.bool_(True)
+    assert_type(int(boolean), int)
+    assert_shape((boolean * 2).shape, ())
+    assert_shape((boolean**2).shape, ())
+    assert_shape((1j * boolean).shape, ())
+    assert boolean - 2.0 == -1.0
+    assert 2.0 - boolean == 1.0
+
+    with assert_raises(TypeError):
+        real + "invalid"  # E: not supported
+
+    with assert_raises(TypeError):
+        boolean - boolean  # E: not supported
+
+    with assert_raises(TypeError):
+        np.float64(1j)  # E: not assignable
+
+
+def test_integer_scalar_array_index() -> None:
+    index = np.ones((), dtype=np.int64)
+    assert_shape(index.shape, ())
+    assert_type(index.__index__(), int)
+    assert [10, 20][index] == 20
+
+    with assert_raises(TypeError):
+        np.ones((), dtype=np.float64).__index__()  # E: not assignable
+
+    with assert_raises(TypeError):
+        np.ones((1,), dtype=np.int64).__index__()  # E: not assignable
+
+
 def test_flexible_and_datetime_scalar_hierarchy() -> None:
     assert_shape(np.zeros((2,), dtype=np.object_).shape, (2,))
     assert_type(np.flexible, type[np.flexible])