diff --git a/src/numpy-stubs/__init__.pyi b/src/numpy-stubs/__init__.pyi index 3b92da3d..abe0a3b9 100644 --- a/src/numpy-stubs/__init__.pyi +++ b/src/numpy-stubs/__init__.pyi @@ -4607,10 +4607,12 @@ class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co] @overload def __init__(self: timedelta64[None], value: _NaTValue | None, format: _TimeUnitSpec, /) -> None: ... @overload - def __init__(self: timedelta64[int], value: dt.timedelta, format: _TimeUnitSpec[_IntTimeUnit], /) -> None: ... + def __init__(self: timedelta64[L[0]], value: L[0], format: _TimeUnitSpec[_IntTD64Unit] = ..., /) -> None: ... @overload def __init__(self: timedelta64[int], value: _IntLike_co, format: _TimeUnitSpec[_IntTD64Unit] = ..., /) -> None: ... @overload + def __init__(self: timedelta64[int], value: dt.timedelta, format: _TimeUnitSpec[_IntTimeUnit], /) -> None: ... + @overload def __init__( self: timedelta64[dt.timedelta], value: dt.timedelta | _IntLike_co, @@ -4645,6 +4647,7 @@ class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co] def __add__(self, x: timedelta64[None], /) -> timedelta64[None]: ... __radd__ = __add__ + # @overload def __mul__(self: timedelta64[_AnyTD64Item], x: int | np.integer | np.bool, /) -> timedelta64[_AnyTD64Item]: ... @overload @@ -4653,30 +4656,73 @@ class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co] def __mul__(self, x: float | np.floating | np.integer | np.bool, /) -> timedelta64: ... __rmul__ = __mul__ + # + @overload + def __mod__(self, x: timedelta64[L[0] | None], /) -> timedelta64[None]: ... @overload def __mod__(self: timedelta64[None], x: timedelta64, /) -> timedelta64[None]: ... @overload + def __mod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int | None]: ... + @overload + def __mod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item | None]: ... + @overload def __mod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> dt.timedelta: ... @overload - def __mod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item]: ... + def __mod__(self, x: timedelta64[int], /) -> timedelta64[int | None]: ... @overload - def __mod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int]: ... + def __mod__(self, x: timedelta64, /) -> timedelta64: ... + + # the L[0] makes __mod__ non-commutative, which the first two overloads reflect @overload - def __mod__(self, x: timedelta64[None], /) -> timedelta64[None]: ... + def __rmod__(self, x: timedelta64[None], /) -> timedelta64[None]: ... @overload - def __mod__(self, x: timedelta64[int], /) -> timedelta64[int]: ... + def __rmod__(self: timedelta64[L[0] | None], x: timedelta64, /) -> timedelta64[None]: ... @overload - def __mod__(self, x: timedelta64, /) -> timedelta64: ... - __rmod__ = __mod__ # at runtime the outcomes differ, but the type signatures are the same + def __rmod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int | None]: ... + @overload + def __rmod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item | None]: ... + @overload + def __rmod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> dt.timedelta: ... + @overload + def __rmod__(self, x: timedelta64[int], /) -> timedelta64[int | None]: ... + @overload + def __rmod__(self, x: timedelta64, /) -> timedelta64: ... + # keep in sync with __mod__ + @overload + def __divmod__(self, x: timedelta64[L[0] | None], /) -> tuple[int64, timedelta64[None]]: ... @overload def __divmod__(self: timedelta64[None], x: timedelta64, /) -> tuple[int64, timedelta64[None]]: ... @overload + def __divmod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> tuple[int64, timedelta64[int | None]]: ... + @overload + def __divmod__( + self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], / + ) -> tuple[int64, timedelta64[_AnyTD64Item | None]]: ... + @overload def __divmod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> tuple[int, dt.timedelta]: ... @overload + def __divmod__(self, x: timedelta64[int], /) -> tuple[int64, timedelta64[int | None]]: ... + @overload def __divmod__(self, x: timedelta64, /) -> tuple[int64, timedelta64]: ... - __rdivmod__ = __divmod__ + # keep in sync with __rmod__ + @overload + def __rdivmod__(self, x: timedelta64[None], /) -> tuple[int64, timedelta64[None]]: ... + @overload + def __rdivmod__(self: timedelta64[L[0] | None], x: timedelta64, /) -> tuple[int64, timedelta64[None]]: ... + @overload + def __rdivmod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> tuple[int64, timedelta64[int | None]]: ... + @overload + def __rdivmod__( + self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], / + ) -> tuple[int64, timedelta64[_AnyTD64Item | None]]: ... + @overload + def __rdivmod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> tuple[int, dt.timedelta]: ... + @overload + def __rdivmod__(self, x: timedelta64[int], /) -> tuple[int64, timedelta64[int | None]]: ... + @overload + def __rdivmod__(self, x: timedelta64, /) -> tuple[int64, timedelta64]: ... @overload def __sub__(self: timedelta64[None], b: _TD64Like_co, /) -> timedelta64[None]: ... @overload @@ -4713,8 +4759,6 @@ class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co] def __truediv__(self: timedelta64[_AnyTD64Item], b: int | integer, /) -> timedelta64[_AnyTD64Item]: ... @overload def __truediv__(self: timedelta64[_AnyTD64Item], b: float | floating, /) -> timedelta64[_AnyTD64Item | None]: ... - - # @overload def __truediv__(self, b: float | floating | integer, /) -> timedelta64: ... @overload @@ -4731,8 +4775,6 @@ class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co] def __floordiv__(self: timedelta64[_AnyTD64Item], b: int | integer, /) -> timedelta64[_AnyTD64Item]: ... @overload def __floordiv__(self: timedelta64[_AnyTD64Item], b: float | floating, /) -> timedelta64[_AnyTD64Item | None]: ... - - # @overload def __rfloordiv__(self: timedelta64[dt.timedelta], a: dt.timedelta, /) -> int: ... @overload diff --git a/test/static/accept/mod.pyi b/test/static/accept/mod.pyi index 45170d84..c214eac5 100644 --- a/test/static/accept/mod.pyi +++ b/test/static/accept/mod.pyi @@ -4,7 +4,7 @@ from typing_extensions import assert_type import numpy as np import numpy.typing as npt -from numpy._typing import _32Bit, _64Bit +from numpy._typing import _64Bit f8: np.float64 i8: np.int64 @@ -14,8 +14,13 @@ f4: np.float32 i4: np.int32 u4: np.uint32 -td: np.timedelta64 -b_: np.bool[L[False]] +m: np.timedelta64 +m_nat: np.timedelta64[None] +m_int0: np.timedelta64[L[0]] +m_int: np.timedelta64[int] +m_td: np.timedelta64[dt.timedelta] + +b_: np.bool b: bool i: int @@ -26,13 +31,40 @@ AR_m: npt.NDArray[np.timedelta64] # Time structures -assert_type(td % td, np.timedelta64[dt.timedelta]) -assert_type(AR_m % td, npt.NDArray[np.timedelta64]) -assert_type(td % AR_m, npt.NDArray[np.timedelta64]) - -assert_type(divmod(td, td), tuple[np.int64, np.timedelta64]) -assert_type(divmod(AR_m, td), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]]) -assert_type(divmod(td, AR_m), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]]) +assert_type(m % m, np.timedelta64) +assert_type(m % m_nat, np.timedelta64[None]) +assert_type(m % m_int0, np.timedelta64[None]) +assert_type(m % m_int, np.timedelta64[int | None]) +assert_type(m_nat % m, np.timedelta64[None]) +assert_type(m_int % m_nat, np.timedelta64[None]) +assert_type(m_int % m_int0, np.timedelta64[None]) +assert_type(m_int % m_int, np.timedelta64[int | None]) +assert_type(m_int % m_td, np.timedelta64[int | None]) +assert_type(m_td % m_nat, np.timedelta64[None]) +assert_type(m_td % m_int0, np.timedelta64[None]) +assert_type(m_td % m_int, np.timedelta64[int | None]) +assert_type(m_td % m_td, np.timedelta64[dt.timedelta | None]) + +assert_type(AR_m % m, npt.NDArray[np.timedelta64]) +assert_type(m % AR_m, npt.NDArray[np.timedelta64]) + +assert_type(divmod(m, m), tuple[np.int64, np.timedelta64]) +assert_type(divmod(m, m_nat), tuple[np.int64, np.timedelta64[None]]) +assert_type(divmod(m, m_int0), tuple[np.int64, np.timedelta64[None]]) +# workarounds for https://github.com/microsoft/pyright/issues/9663 +assert_type(m.__divmod__(m_int), tuple[np.int64, np.timedelta64[int | None]]) +assert_type(divmod(m_nat, m), tuple[np.int64, np.timedelta64[None]]) +assert_type(divmod(m_int, m_nat), tuple[np.int64, np.timedelta64[None]]) +assert_type(divmod(m_int, m_int0), tuple[np.int64, np.timedelta64[None]]) +assert_type(divmod(m_int, m_int), tuple[np.int64, np.timedelta64[int | None]]) +assert_type(divmod(m_int, m_td), tuple[np.int64, np.timedelta64[int | None]]) +assert_type(divmod(m_td, m_nat), tuple[np.int64, np.timedelta64[None]]) +assert_type(divmod(m_td, m_int0), tuple[np.int64, np.timedelta64[None]]) +assert_type(divmod(m_td, m_int), tuple[np.int64, np.timedelta64[int | None]]) +assert_type(divmod(m_td, m_td), tuple[np.int64, np.timedelta64[dt.timedelta | None]]) + +assert_type(divmod(AR_m, m), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]]) +assert_type(divmod(m, AR_m), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]]) # Bool @@ -46,11 +78,12 @@ assert_type(b_ % f8, np.float64) assert_type(b_ % AR_b, npt.NDArray[np.int8]) assert_type(divmod(b_, b), tuple[np.int8, np.int8]) -assert_type(divmod(b_, i), tuple[np.int_, np.int_]) -assert_type(divmod(b_, f), tuple[np.float64, np.float64]) assert_type(divmod(b_, b_), tuple[np.int8, np.int8]) -assert_type(divmod(b_, i8), tuple[np.int64, np.int64]) -assert_type(divmod(b_, u8), tuple[np.uint64, np.uint64]) +# workarounds for https://github.com/microsoft/pyright/issues/9663 +assert_type(b_.__divmod__(i), tuple[np.int_, np.int_]) +assert_type(b_.__divmod__(f), tuple[np.float64, np.float64]) +assert_type(b_.__divmod__(i8), tuple[np.int64, np.int64]) +assert_type(b_.__divmod__(u8), tuple[np.uint64, np.uint64]) assert_type(divmod(b_, f8), tuple[np.float64, np.float64]) assert_type(divmod(b_, AR_b), tuple[npt.NDArray[np.int8], npt.NDArray[np.int8]]) @@ -76,29 +109,27 @@ assert_type(divmod(AR_b, b_), tuple[npt.NDArray[np.int8], npt.NDArray[np.int8]]) assert_type(i8 % b, np.int64) assert_type(i8 % i8, np.int64) -assert_type(i8 % f, np.floating[_64Bit]) -assert_type(i8 % f8, np.floating[_64Bit]) +assert_type(i8 % f, np.float64 | np.floating[_64Bit]) +assert_type(i8 % f8, np.float64 | np.floating[_64Bit]) assert_type(i4 % i8, np.int64 | np.int32) assert_type(i4 % f8, np.float64 | np.float32) assert_type(i4 % i4, np.int32) assert_type(i4 % f4, np.float32) assert_type(i8 % AR_b, npt.NDArray[np.int64]) -assert_type(divmod(i8, b), tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]]) -assert_type(divmod(i8, f), tuple[np.floating[_64Bit], np.floating[_64Bit]]) -assert_type(divmod(i8, i8), tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]]) -assert_type(divmod(i8, f8), tuple[np.floating[_64Bit], np.floating[_64Bit]]) -assert_type( - divmod(i8, i4), - tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]] | tuple[np.signedinteger[_32Bit], np.signedinteger[_32Bit]], -) -assert_type(divmod(i8, f4), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.floating[_32Bit], np.floating[_32Bit]]) -assert_type(divmod(i4, i4), tuple[np.signedinteger[_32Bit], np.signedinteger[_32Bit]]) -assert_type(divmod(i4, f4), tuple[np.floating[_32Bit], np.floating[_32Bit]]) +assert_type(divmod(i8, b), tuple[np.int64, np.int64]) +assert_type(divmod(i8, i4), tuple[np.int64, np.int64] | tuple[np.int32, np.int32]) +assert_type(divmod(i8, i8), tuple[np.int64, np.int64]) +# workarounds for https://github.com/microsoft/pyright/issues/9663 +assert_type(i8.__divmod__(f), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64]) +assert_type(i8.__divmod__(f8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64]) +assert_type(divmod(i8, f4), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float32, np.float32]) +assert_type(divmod(i4, i4), tuple[np.int32, np.int32]) +assert_type(divmod(i4, f4), tuple[np.float32, np.float32]) assert_type(divmod(i8, AR_b), tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]) -assert_type(b % i8, np.signedinteger[_64Bit]) -assert_type(f % i8, np.floating[_64Bit]) +assert_type(b % i8, np.int64) +assert_type(f % i8, np.float64 | np.floating[_64Bit]) assert_type(i8 % i8, np.int64) assert_type(f8 % i8, np.float64) assert_type(i8 % i4, np.int64 | np.int32) @@ -107,24 +138,22 @@ assert_type(i4 % i4, np.int32) assert_type(f4 % i4, np.float32) assert_type(AR_b % i8, npt.NDArray[np.int64]) -assert_type(divmod(b, i8), tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]]) -assert_type(divmod(f, i8), tuple[np.floating[_64Bit], np.floating[_64Bit]]) +assert_type(divmod(b, i8), tuple[np.int64, np.int64]) +assert_type(divmod(f, i8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64]) assert_type(divmod(i8, i8), tuple[np.int64, np.int64]) assert_type(divmod(f8, i8), tuple[np.float64, np.float64]) -assert_type( - divmod(i4, i8), - tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]] | tuple[np.signedinteger[_32Bit], np.signedinteger[_32Bit]], -) -assert_type(divmod(f4, i8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.floating[_32Bit], np.floating[_32Bit]]) -assert_type(divmod(i4, i4), tuple[np.signedinteger[_32Bit], np.signedinteger[_32Bit]]) -assert_type(divmod(f4, i4), tuple[np.floating[_32Bit], np.floating[_32Bit]]) -assert_type(divmod(AR_b, i8), tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]) +assert_type(divmod(i4, i8), tuple[np.int64, np.int64] | tuple[np.int32, np.int32]) +assert_type(divmod(i4, i4), tuple[np.int32, np.int32]) +# workarounds for https://github.com/microsoft/pyright/issues/9663 +assert_type(f4.__divmod__(i8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float32, np.float32]) +assert_type(f4.__divmod__(i4), tuple[np.float32, np.float32]) +assert_type(AR_b.__divmod__(i8), tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]) # float assert_type(f8 % b, np.float64) assert_type(f8 % f, np.float64) -assert_type(i8 % f4, np.floating[_64Bit] | np.floating[_32Bit]) +assert_type(i8 % f4, np.floating[_64Bit] | np.float32) assert_type(f4 % f4, np.float32) assert_type(f8 % AR_b, npt.NDArray[np.float64]) @@ -136,15 +165,16 @@ assert_type(divmod(f4, f4), tuple[np.float32, np.float32]) assert_type(divmod(f8, AR_b), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]) assert_type(b % f8, np.float64) -assert_type(f % f8, np.float64) +assert_type(f % f8, np.float64) # pyright: ignore[reportAssertTypeFailure] # pyright incorrectly infers `builtins.float` assert_type(f8 % f8, np.float64) assert_type(f8 % f8, np.float64) assert_type(f4 % f4, np.float32) assert_type(AR_b % f8, npt.NDArray[np.float64]) assert_type(divmod(b, f8), tuple[np.float64, np.float64]) -assert_type(divmod(f, f8), tuple[np.float64, np.float64]) assert_type(divmod(f8, f8), tuple[np.float64, np.float64]) -assert_type(divmod(f4, f8), tuple[np.float64, np.float64] | tuple[np.float32, np.float32]) assert_type(divmod(f4, f4), tuple[np.float32, np.float32]) -assert_type(divmod(AR_b, f8), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]) +# workarounds for https://github.com/microsoft/pyright/issues/9663 +assert_type(f8.__rdivmod__(f), tuple[np.float64, np.float64]) +assert_type(f8.__rdivmod__(f4), tuple[np.float64, np.float64]) +assert_type(AR_b.__divmod__(f8), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])