From e98a9407b94e6dc2caf9714aa57d9ba3479fc1e5 Mon Sep 17 00:00:00 2001 From: jorenham Date: Mon, 18 Nov 2024 18:53:36 +0100 Subject: [PATCH] TYP: Generic ``timedelta64`` and ``datetime64`` scalar types --- numpy/__init__.pyi | 276 +++++++++++++++--- numpy/_typing/_callable.pyi | 11 - numpy/typing/tests/data/reveal/arithmetic.pyi | 131 ++++++--- numpy/typing/tests/data/reveal/mod.pyi | 3 +- 4 files changed, 327 insertions(+), 94 deletions(-) diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 3410cec43450..57c2b9b8af39 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -143,7 +143,6 @@ from numpy._typing._callable import ( _BoolTrueDiv, _BoolMod, _BoolDivMod, - _TD64Div, _IntTrueDiv, _UnsignedIntOp, _UnsignedIntBitOp, @@ -766,6 +765,11 @@ _AnyShapeType = TypeVar( tuple[int, ...], # N-d ) _AnyNBitInexact = TypeVar("_AnyNBitInexact", _NBitHalf, _NBitSingle, _NBitDouble, _NBitLongDouble) +_AnyTD64Item = TypeVar("_AnyTD64Item", dt.timedelta, int, None, dt.timedelta | int | None) +_AnyDT64Arg = TypeVar("_AnyDT64Arg", dt.datetime, dt.date, None) +_AnyDT64Item = TypeVar("_AnyDT64Item", dt.datetime, dt.date, int, None, dt.date, int | None) +_AnyDate = TypeVar("_AnyDate", dt.date, dt.datetime) +_AnyDateOrTime = TypeVar("_AnyDateOrTime", dt.date, dt.datetime, dt.timedelta) ### Type parameters (for internal use only) @@ -813,6 +817,9 @@ _FlexibleItemT_co = TypeVar( covariant=True, ) _CharacterItemT_co = TypeVar("_CharacterItemT_co", bound=_CharLike_co, default=_CharLike_co, covariant=True) +_TD64ItemT_co = TypeVar("_TD64ItemT_co", bound=dt.timedelta | int | None, default=dt.timedelta | int | None, covariant=True) +_DT64ItemT_co = TypeVar("_DT64ItemT_co", bound=dt.date | int | None, default=dt.date | int | None, covariant=True) +_TD64UnitT = TypeVar("_TD64UnitT", bound=_TD64Unit, default=_TD64Unit) ### Type Aliases (for internal use only) @@ -1010,6 +1017,21 @@ _MemMapModeKind: TypeAlias = L[ "write", "w+", ] +_DT64Date: TypeAlias = _HasDateAttributes | L["TODAY", "today", b"TODAY", b"today"] +_DT64Now: TypeAlias = L["NOW", "now", b"NOW", b"now"] +_NaTValue: TypeAlias = L["NAT","NaT", "nat",b"NAT", b"NaT", b"nat"] + +_MonthUnit: TypeAlias = L["Y", "M", b"Y", b"M"] +_DayUnit: TypeAlias = L["W", "D", b"W", b"D"] +_DateUnit: TypeAlias = L[_MonthUnit, _DayUnit] +_NativeTimeUnit: TypeAlias = L["h", "m", "s", "ms", "us", "μs", b"h", b"m", b"s", b"ms", b"us"] +_IntTimeUnit: TypeAlias = L["ns", "ps", "fs", "as", b"ns", b"ps", b"fs", b"as"] +_TimeUnit: TypeAlias = L[_NativeTimeUnit, _IntTimeUnit] +_NativeTD64Unit: TypeAlias = L[_DayUnit, _NativeTimeUnit] +_IntTD64Unit: TypeAlias = L[_MonthUnit, _IntTimeUnit] +_TD64Unit: TypeAlias = L[_DateUnit, _TimeUnit] +_TimeUnitSpec: TypeAlias = _TD64UnitT | tuple[_TD64UnitT, SupportsIndex] + ### Protocols (for internal use only) @type_check_only @@ -4010,68 +4032,238 @@ csingle: TypeAlias = complexfloating[_NBitSingle, _NBitSingle] cdouble: TypeAlias = complexfloating[_NBitDouble, _NBitDouble] clongdouble: TypeAlias = complexfloating[_NBitLongDouble, _NBitLongDouble] -class timedelta64(_IntegralMixin, generic[dt.timedelta | int | None]): - def __init__( - self, - value: timedelta64 | dt.timedelta | int | _CharLike_co | None = None, - format: _CharLike_co | tuple[_CharLike_co, _IntLike_co] = ..., - /, - ) -> None: ... - +class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co]): @property def itemsize(self) -> L[8]: ... @property def nbytes(self) -> L[8]: ... + @overload + def __init__(self, value: _TD64ItemT_co | timedelta64[_TD64ItemT_co], /) -> None: ... + @overload + def __init__(self: timedelta64[L[0]], /) -> None: ... + @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: ... + @overload + def __init__(self: timedelta64[int], value: int, format: _TimeUnitSpec[_IntTD64Unit] = ..., /) -> None: ... + @overload + def __init__( + self: timedelta64[dt.timedelta], + value: dt.timedelta | int, + format: _TimeUnitSpec[_NativeTD64Unit] = ..., + /, + ) -> None: ... + @overload + def __init__(self, value: int | bytes | str | dt.timedelta | None, format: _TimeUnitSpec = ..., /) -> None: ... + # NOTE: Only a limited number of units support conversion # to builtin scalar types: `Y`, `M`, `ns`, `ps`, `fs`, `as` - def __neg__(self) -> Self: ... - def __pos__(self) -> Self: ... - def __abs__(self) -> Self: ... + def __int__(self: timedelta64[int], /) -> int: ... + def __float__(self: timedelta64[int], /) -> float: ... - def __add__(self, other: _TD64Like_co, /) -> timedelta64: ... - def __radd__(self, other: _TD64Like_co, /) -> timedelta64: ... - def __sub__(self, other: _TD64Like_co, /) -> timedelta64: ... - def __rsub__(self, other: _TD64Like_co, /) -> timedelta64: ... - def __mul__(self, other: _FloatLike_co, /) -> timedelta64: ... - def __rmul__(self, other: _FloatLike_co, /) -> timedelta64: ... - __truediv__: _TD64Div[float64] - def __rtruediv__(self, other: timedelta64, /) -> float64: ... - __floordiv__: _TD64Div[int64] - def __rfloordiv__(self, other: timedelta64, /) -> int64: ... - def __mod__(self, other: timedelta64, /) -> timedelta64: ... - def __rmod__(self, other: timedelta64, /) -> timedelta64: ... - def __divmod__(self, other: timedelta64, /) -> tuple[int64, timedelta64]: ... - def __rdivmod__(self, other: timedelta64, /) -> tuple[int64, timedelta64]: ... + def __neg__(self, /) -> Self: ... + def __pos__(self, /) -> Self: ... + def __abs__(self, /) -> Self: ... + + @overload + def __add__(self: timedelta64[None], x: _TD64Like_co, /) -> timedelta64[None]: ... + @overload + def __add__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int]: ... + @overload + def __add__(self: timedelta64[int], x: timedelta64, /) -> timedelta64[int | None]: ... + @overload + def __add__(self: timedelta64[dt.timedelta], x: _AnyDateOrTime, /) -> _AnyDateOrTime: ... + @overload + def __add__(self: timedelta64[_AnyTD64Item], x: timedelta64[_AnyTD64Item] | _IntLike_co, /) -> timedelta64[_AnyTD64Item]: ... + @overload + def __add__(self, x: timedelta64[None], /) -> timedelta64[None]: ... + __radd__ = __add__ + + @overload + def __mul__(self: timedelta64[_AnyTD64Item], x: int | np.integer[Any] | np.bool, /) -> timedelta64[_AnyTD64Item]: ... + @overload + def __mul__(self: timedelta64[_AnyTD64Item], x: float | np.floating[Any], /) -> timedelta64[_AnyTD64Item | None]: ... + @overload + def __mul__(self, x: float | np.floating[Any] | np.integer[Any] | np.bool, /) -> timedelta64: ... + __rmul__ = __mul__ + + @overload + def __mod__(self: timedelta64[None], x: timedelta64, /) -> timedelta64[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]: ... + @overload + def __mod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int]: ... + @overload + def __mod__(self, x: timedelta64[None], /) -> timedelta64[None]: ... + @overload + def __mod__(self, x: timedelta64[int], /) -> timedelta64[int]: ... + @overload + def __mod__(self, x: timedelta64, /) -> timedelta64: ... + __rmod__ = __mod__ # at runtime the outcomes differ, but the type signatures are the same + + @overload + def __divmod__(self: timedelta64[None], x: timedelta64, /) -> tuple[int64, timedelta64[None]]: ... + @overload + def __divmod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> tuple[int, dt.timedelta]: ... + @overload + def __divmod__(self, x: timedelta64, /) -> tuple[int64, timedelta64]: ... + __rdivmod__ = __divmod__ + + @overload + def __sub__(self: timedelta64[None], b: _TD64Like_co, /) -> timedelta64[None]: ... + @overload + def __sub__(self: timedelta64[int], b: timedelta64[int | dt.timedelta], /) -> timedelta64[int]: ... + @overload + def __sub__(self: timedelta64[int], b: timedelta64, /) -> timedelta64[int | None]: ... + @overload + def __sub__(self: timedelta64[dt.timedelta], b: dt.timedelta, /) -> dt.timedelta: ... + @overload + def __sub__(self: timedelta64[_AnyTD64Item], b: timedelta64[_AnyTD64Item] | _IntLike_co, /) -> timedelta64[_AnyTD64Item]: ... + @overload + def __sub__(self, b: timedelta64[None], /) -> timedelta64[None]: ... + + @overload + def __rsub__(self: timedelta64[None], a: _TD64Like_co, /) -> timedelta64[None]: ... + @overload + def __rsub__(self: timedelta64[dt.timedelta], a: _AnyDateOrTime, /) -> _AnyDateOrTime: ... + @overload + def __rsub__(self: timedelta64[dt.timedelta], a: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item]: ... + @overload + def __rsub__(self: timedelta64[_AnyTD64Item], a: timedelta64[_AnyTD64Item] | _IntLike_co, /) -> timedelta64[_AnyTD64Item]: ... + @overload + def __rsub__(self, a: timedelta64[None], /) -> timedelta64[None]: ... + @overload + def __rsub__(self, a: datetime64[None], /) -> datetime64[None]: ... + + @overload + def __truediv__(self: timedelta64[dt.timedelta], b: dt.timedelta, /) -> float: ... + @overload + def __truediv__(self, b: timedelta64, /) -> float64: ... + @overload + 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 + def __rtruediv__(self: timedelta64[dt.timedelta], a: dt.timedelta, /) -> float: ... + @overload + def __rtruediv__(self, a: timedelta64, /) -> float64: ... + + @overload + def __floordiv__(self: timedelta64[dt.timedelta], b: dt.timedelta, /) -> int: ... + @overload + def __floordiv__(self, b: timedelta64, /) -> int64: ... + @overload + 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 + def __rfloordiv__(self, a: timedelta64, /) -> int64: ... __lt__: _ComparisonOpLT[_TD64Like_co, _ArrayLikeTD64_co] __le__: _ComparisonOpLE[_TD64Like_co, _ArrayLikeTD64_co] __gt__: _ComparisonOpGT[_TD64Like_co, _ArrayLikeTD64_co] __ge__: _ComparisonOpGE[_TD64Like_co, _ArrayLikeTD64_co] -class datetime64(_RealMixin, generic[dt.date | int | None]): - @overload - def __init__(self, value: int, format: _CharLike_co | tuple[_CharLike_co, _IntLike_co], /) -> None: ... - @overload - def __init__( - self, - value: datetime64 | _CharLike_co | _HasDateAttributes | None = None, - format: _CharLike_co | tuple[_CharLike_co, _IntLike_co] = ..., - /, - ) -> None: ... - +class datetime64(_RealMixin, generic[_DT64ItemT_co], Generic[_DT64ItemT_co]): @property def itemsize(self) -> L[8]: ... @property def nbytes(self) -> L[8]: ... - def __add__(self, other: _TD64Like_co, /) -> datetime64: ... - def __radd__(self, other: _TD64Like_co, /) -> datetime64: ... @overload - def __sub__(self, other: datetime64, /) -> timedelta64: ... + def __init__(self, value: datetime64[_DT64ItemT_co], /) -> None: ... + @overload + def __init__(self: datetime64[_AnyDT64Arg], value: _AnyDT64Arg, /) -> None: ... + @overload + def __init__(self: datetime64[None], value: _NaTValue | None = ..., format: _TimeUnitSpec = ..., /) -> None: ... + @overload + def __init__(self: datetime64[dt.datetime], value: _DT64Now, format: _TimeUnitSpec[_NativeTimeUnit] = ..., /) -> None: ... + @overload + def __init__(self: datetime64[dt.date], value: _DT64Date, format: _TimeUnitSpec[_DateUnit] = ..., /) -> None: ... + @overload + def __init__(self: datetime64[int], value: int | bytes | str | dt.date, format: _TimeUnitSpec[_IntTimeUnit], /) -> None: ... + @overload + def __init__(self: datetime64[dt.datetime], value: int | bytes | str, format: _TimeUnitSpec[_NativeTimeUnit], /) -> None: ... + @overload + def __init__(self: datetime64[dt.date], value: int | bytes | str, format: _TimeUnitSpec[_DateUnit], /) -> None: ... + @overload + def __init__(self, value: bytes | str | None, format: _TimeUnitSpec = ..., /) -> None: ... + + @overload + def __add__(self: datetime64[_AnyDT64Item], x: int | integer[Any] | np.bool, /) -> datetime64[_AnyDT64Item]: ... + @overload + def __add__(self: datetime64[None], x: _TD64Like_co, /) -> datetime64[None]: ... + @overload + def __add__(self: datetime64[int], x: timedelta64[int | dt.timedelta], /) -> datetime64[int]: ... + @overload + def __add__(self: datetime64[dt.datetime], x: timedelta64[dt.timedelta], /) -> datetime64[dt.datetime]: ... + @overload + def __add__(self: datetime64[dt.date], x: timedelta64[dt.timedelta], /) -> datetime64[dt.date]: ... + @overload + def __add__(self: datetime64[dt.date], x: timedelta64[int], /) -> datetime64[int]: ... + @overload + def __add__(self, x: datetime64[None], /) -> datetime64[None]: ... + @overload + def __add__(self, x: _TD64Like_co, /) -> datetime64: ... + __radd__ = __add__ + + @overload + def __sub__(self: datetime64[_AnyDT64Item], x: int | integer[Any] | np.bool, /) -> datetime64[_AnyDT64Item]: ... + @overload + def __sub__(self: datetime64[_AnyDate], x: _AnyDate, /) -> dt.timedelta: ... + @overload + def __sub__(self: datetime64[None], x: timedelta64, /) -> datetime64[None]: ... + @overload + def __sub__(self: datetime64[None], x: datetime64, /) -> timedelta64[None]: ... + @overload + def __sub__(self: datetime64[int], x: timedelta64, /) -> datetime64[int]: ... + @overload + def __sub__(self: datetime64[int], x: datetime64, /) -> timedelta64[int]: ... + @overload + def __sub__(self: datetime64[dt.datetime], x: timedelta64[int], /) -> datetime64[int]: ... + @overload + def __sub__(self: datetime64[dt.datetime], x: timedelta64[dt.timedelta], /) -> datetime64[dt.datetime]: ... + @overload + def __sub__(self: datetime64[dt.datetime], x: datetime64[int], /) -> timedelta64[int]: ... + @overload + def __sub__(self: datetime64[dt.date], x: timedelta64[int], /) -> datetime64[dt.date | int]: ... + @overload + def __sub__(self: datetime64[dt.date], x: timedelta64[dt.timedelta], /) -> datetime64[dt.date]: ... + @overload + def __sub__(self: datetime64[dt.date], x: datetime64[dt.date], /) -> timedelta64[dt.timedelta]: ... + @overload + def __sub__(self, x: timedelta64[None], /) -> datetime64[None]: ... + @overload + def __sub__(self, x: datetime64[None], /) -> timedelta64[None]: ... + @overload + def __sub__(self, x: _TD64Like_co, /) -> datetime64: ... + @overload + def __sub__(self, x: datetime64, /) -> timedelta64: ... + + @overload + def __rsub__(self: datetime64[_AnyDT64Item], x: int | integer[Any] | np.bool, /) -> datetime64[_AnyDT64Item]: ... + @overload + def __rsub__(self: datetime64[_AnyDate], x: _AnyDate, /) -> dt.timedelta: ... + @overload + def __rsub__(self: datetime64[None], x: datetime64, /) -> timedelta64[None]: ... + @overload + def __rsub__(self: datetime64[int], x: datetime64, /) -> timedelta64[int]: ... + @overload + def __rsub__(self: datetime64[dt.datetime], x: datetime64[int], /) -> timedelta64[int]: ... + @overload + def __rsub__(self: datetime64[dt.datetime], x: datetime64[dt.date], /) -> timedelta64[dt.timedelta]: ... + @overload + def __rsub__(self, x: datetime64[None], /) -> timedelta64[None]: ... @overload - def __sub__(self, other: _TD64Like_co, /) -> datetime64: ... - def __rsub__(self, other: datetime64, /) -> timedelta64: ... + def __rsub__(self, x: datetime64, /) -> timedelta64: ... __lt__: _ComparisonOpLT[datetime64, _ArrayLikeDT64_co] __le__: _ComparisonOpLE[datetime64, _ArrayLikeDT64_co] diff --git a/numpy/_typing/_callable.pyi b/numpy/_typing/_callable.pyi index ec82d2a923d1..56e24fb73911 100644 --- a/numpy/_typing/_callable.pyi +++ b/numpy/_typing/_callable.pyi @@ -22,7 +22,6 @@ from typing import ( import numpy as np from numpy import ( generic, - timedelta64, number, integer, unsignedinteger, @@ -38,7 +37,6 @@ from ._nbit import _NBitInt from ._scalars import ( _BoolLike_co, _IntLike_co, - _FloatLike_co, _NumberLike_co, ) from . import NBitBase @@ -132,15 +130,6 @@ class _BoolDivMod(Protocol): @overload def __call__(self, other: _FloatType, /) -> _2Tuple[_FloatType]: ... -@type_check_only -class _TD64Div(Protocol[_NumberType_co]): - @overload - def __call__(self, other: timedelta64, /) -> _NumberType_co: ... - @overload - def __call__(self, other: _BoolLike_co, /) -> NoReturn: ... - @overload - def __call__(self, other: _FloatLike_co, /) -> timedelta64: ... - @type_check_only class _IntTrueDiv(Protocol[_NBit1]): @overload diff --git a/numpy/typing/tests/data/reveal/arithmetic.pyi b/numpy/typing/tests/data/reveal/arithmetic.pyi index 1adfb2fcfce8..c1eee5d3fc29 100644 --- a/numpy/typing/tests/data/reveal/arithmetic.pyi +++ b/numpy/typing/tests/data/reveal/arithmetic.pyi @@ -1,3 +1,4 @@ +import datetime as dt from typing import Any import numpy as np @@ -6,28 +7,40 @@ from numpy._typing import _32Bit,_64Bit, _128Bit from typing_extensions import assert_type +b: bool +c: complex +f: float +i: int + +c16: np.complex128 +c8: np.complex64 + # Can't directly import `np.float128` as it is not available on all platforms f16: np.floating[_128Bit] +f8: np.float64 +f4: np.float32 -c16 = np.complex128() -f8 = np.float64() -i8 = np.int64() -u8 = np.uint64() +i8: np.int64 +i4: np.int32 -c8 = np.complex64() -f4 = np.float32() -i4 = np.int32() -u4 = np.uint32() +u8: np.uint64 +u4: np.uint32 -dt = np.datetime64(0, "D") -td = np.timedelta64(0, "D") +b_: np.bool -b_ = np.bool() +M8: np.datetime64 +M8_none: np.datetime64[None] +M8_date: np.datetime64[dt.date] +M8_time: np.datetime64[dt.datetime] +M8_int: np.datetime64[int] +date: dt.date +time: dt.datetime -b = bool() -c = complex() -f = float() -i = int() +m8: np.timedelta64 +m8_none: np.timedelta64[None] +m8_int: np.timedelta64[int] +m8_delta: np.timedelta64[dt.timedelta] +delta: dt.timedelta AR_b: npt.NDArray[np.bool] AR_u: npt.NDArray[np.uint32] @@ -259,7 +272,10 @@ assert_type(-i8, np.int64) assert_type(-i4, np.int32) assert_type(-u8, np.uint64) assert_type(-u4, np.uint32) -assert_type(-td, np.timedelta64) +assert_type(-m8, np.timedelta64) +assert_type(-m8_none, np.timedelta64[None]) +assert_type(-m8_int, np.timedelta64[int]) +assert_type(-m8_delta, np.timedelta64[dt.timedelta]) assert_type(-AR_f, npt.NDArray[np.float64]) assert_type(+f16, np.floating[_128Bit]) @@ -271,7 +287,9 @@ assert_type(+i8, np.int64) assert_type(+i4, np.int32) assert_type(+u8, np.uint64) assert_type(+u4, np.uint32) -assert_type(+td, np.timedelta64) +assert_type(+m8_none, np.timedelta64[None]) +assert_type(+m8_int, np.timedelta64[int]) +assert_type(+m8_delta, np.timedelta64[dt.timedelta]) assert_type(+AR_f, npt.NDArray[np.float64]) assert_type(abs(f16), np.floating[_128Bit]) @@ -283,33 +301,66 @@ assert_type(abs(i8), np.int64) assert_type(abs(i4), np.int32) assert_type(abs(u8), np.uint64) assert_type(abs(u4), np.uint32) -assert_type(abs(td), np.timedelta64) +assert_type(abs(m8), np.timedelta64) +assert_type(abs(m8_none), np.timedelta64[None]) +assert_type(abs(m8_int), np.timedelta64[int]) +assert_type(abs(m8_delta), np.timedelta64[dt.timedelta]) assert_type(abs(b_), np.bool) # Time structures -assert_type(dt + td, np.datetime64) -assert_type(dt + i, np.datetime64) -assert_type(dt + i4, np.datetime64) -assert_type(dt + i8, np.datetime64) -assert_type(dt - dt, np.timedelta64) -assert_type(dt - i, np.datetime64) -assert_type(dt - i4, np.datetime64) -assert_type(dt - i8, np.datetime64) - -assert_type(td + td, np.timedelta64) -assert_type(td + i, np.timedelta64) -assert_type(td + i4, np.timedelta64) -assert_type(td + i8, np.timedelta64) -assert_type(td - td, np.timedelta64) -assert_type(td - i, np.timedelta64) -assert_type(td - i4, np.timedelta64) -assert_type(td - i8, np.timedelta64) -assert_type(td / f, np.timedelta64) -assert_type(td / f4, np.timedelta64) -assert_type(td / f8, np.timedelta64) -assert_type(td / td, np.float64) -assert_type(td // td, np.int64) +assert_type(M8 + m8, np.datetime64) +assert_type(M8 + i, np.datetime64) +assert_type(M8 + i8, np.datetime64) +assert_type(M8 - M8, np.timedelta64) +assert_type(M8 - i, np.datetime64) +assert_type(M8 - i8, np.datetime64) + +assert_type(M8_none + m8, np.datetime64[None]) +assert_type(M8_none + i, np.datetime64[None]) +assert_type(M8_none + i8, np.datetime64[None]) +assert_type(M8_none - M8, np.timedelta64[None]) +assert_type(M8_none - m8, np.datetime64[None]) +assert_type(M8_none - i, np.datetime64[None]) +assert_type(M8_none - i8, np.datetime64[None]) + +assert_type(m8 + m8, np.timedelta64) +assert_type(m8 + i, np.timedelta64) +assert_type(m8 + i8, np.timedelta64) +assert_type(m8 - m8, np.timedelta64) +assert_type(m8 - i, np.timedelta64) +assert_type(m8 - i8, np.timedelta64) +assert_type(m8 * f, np.timedelta64) +assert_type(m8 * f4, np.timedelta64) +assert_type(m8 * np.True_, np.timedelta64) +assert_type(m8 / f, np.timedelta64) +assert_type(m8 / f4, np.timedelta64) +assert_type(m8 / m8, np.float64) +assert_type(m8 // m8, np.int64) +assert_type(m8 % m8, np.timedelta64) +assert_type(divmod(m8, m8), tuple[np.int64, np.timedelta64]) + +assert_type(m8_none + m8, np.timedelta64[None]) +assert_type(m8_none + i, np.timedelta64[None]) +assert_type(m8_none + i8, np.timedelta64[None]) +assert_type(m8_none - i, np.timedelta64[None]) +assert_type(m8_none - i8, np.timedelta64[None]) + +assert_type(m8_int + i, np.timedelta64[int]) +assert_type(m8_int + m8_delta, np.timedelta64[int]) +assert_type(m8_int + m8, np.timedelta64[int | None]) +assert_type(m8_int - i, np.timedelta64[int]) +assert_type(m8_int - m8_delta, np.timedelta64[int]) +assert_type(m8_int - m8, np.timedelta64[int | None]) + +assert_type(m8_delta + date, dt.date) +assert_type(m8_delta + time, dt.datetime) +assert_type(m8_delta + delta, dt.timedelta) +assert_type(m8_delta - delta, dt.timedelta) +assert_type(m8_delta / delta, float) +assert_type(m8_delta // delta, int) +assert_type(m8_delta % delta, dt.timedelta) +assert_type(divmod(m8_delta, delta), tuple[int, dt.timedelta]) # boolean diff --git a/numpy/typing/tests/data/reveal/mod.pyi b/numpy/typing/tests/data/reveal/mod.pyi index b0b3864351a6..e7e6082753be 100644 --- a/numpy/typing/tests/data/reveal/mod.pyi +++ b/numpy/typing/tests/data/reveal/mod.pyi @@ -1,3 +1,4 @@ +import datetime as dt from typing import Any import numpy as np @@ -26,7 +27,7 @@ AR_m: npt.NDArray[np.timedelta64] # Time structures -assert_type(td % td, np.timedelta64) +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])