Skip to content

Commit 484f9bf

Browse files
charrischarris
authored
Merge pull request numpy#27790 from jorenham/typing/generic-time-scalars
TYP: Generic ``timedelta64`` and ``datetime64`` scalar types
2 parents f3cca78 + e98a940 commit 484f9bf

File tree

4 files changed

+327
-94
lines changed

4 files changed

+327
-94
lines changed

‎numpy/__init__.pyi

Lines changed: 234 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,6 @@ from numpy._typing._callable import (
143143
_BoolTrueDiv,
144144
_BoolMod,
145145
_BoolDivMod,
146-
_TD64Div,
147146
_IntTrueDiv,
148147
_UnsignedIntOp,
149148
_UnsignedIntBitOp,
@@ -766,6 +765,11 @@ _AnyShapeType = TypeVar(
766765
tuple[int, ...], # N-d
767766
)
768767
_AnyNBitInexact = TypeVar("_AnyNBitInexact", _NBitHalf, _NBitSingle, _NBitDouble, _NBitLongDouble)
768+
_AnyTD64Item = TypeVar("_AnyTD64Item", dt.timedelta, int, None, dt.timedelta | int | None)
769+
_AnyDT64Arg = TypeVar("_AnyDT64Arg", dt.datetime, dt.date, None)
770+
_AnyDT64Item = TypeVar("_AnyDT64Item", dt.datetime, dt.date, int, None, dt.date, int | None)
771+
_AnyDate = TypeVar("_AnyDate", dt.date, dt.datetime)
772+
_AnyDateOrTime = TypeVar("_AnyDateOrTime", dt.date, dt.datetime, dt.timedelta)
769773

770774
### Type parameters (for internal use only)
771775

@@ -813,6 +817,9 @@ _FlexibleItemT_co = TypeVar(
813817
covariant=True,
814818
)
815819
_CharacterItemT_co = TypeVar("_CharacterItemT_co", bound=_CharLike_co, default=_CharLike_co, covariant=True)
820+
_TD64ItemT_co = TypeVar("_TD64ItemT_co", bound=dt.timedelta | int | None, default=dt.timedelta | int | None, covariant=True)
821+
_DT64ItemT_co = TypeVar("_DT64ItemT_co", bound=dt.date | int | None, default=dt.date | int | None, covariant=True)
822+
_TD64UnitT = TypeVar("_TD64UnitT", bound=_TD64Unit, default=_TD64Unit)
816823

817824
### Type Aliases (for internal use only)
818825

@@ -1010,6 +1017,21 @@ _MemMapModeKind: TypeAlias = L[
10101017
"write", "w+",
10111018
]
10121019

1020+
_DT64Date: TypeAlias = _HasDateAttributes | L["TODAY", "today", b"TODAY", b"today"]
1021+
_DT64Now: TypeAlias = L["NOW", "now", b"NOW", b"now"]
1022+
_NaTValue: TypeAlias = L["NAT","NaT", "nat",b"NAT", b"NaT", b"nat"]
1023+
1024+
_MonthUnit: TypeAlias = L["Y", "M", b"Y", b"M"]
1025+
_DayUnit: TypeAlias = L["W", "D", b"W", b"D"]
1026+
_DateUnit: TypeAlias = L[_MonthUnit, _DayUnit]
1027+
_NativeTimeUnit: TypeAlias = L["h", "m", "s", "ms", "us", "μs", b"h", b"m", b"s", b"ms", b"us"]
1028+
_IntTimeUnit: TypeAlias = L["ns", "ps", "fs", "as", b"ns", b"ps", b"fs", b"as"]
1029+
_TimeUnit: TypeAlias = L[_NativeTimeUnit, _IntTimeUnit]
1030+
_NativeTD64Unit: TypeAlias = L[_DayUnit, _NativeTimeUnit]
1031+
_IntTD64Unit: TypeAlias = L[_MonthUnit, _IntTimeUnit]
1032+
_TD64Unit: TypeAlias = L[_DateUnit, _TimeUnit]
1033+
_TimeUnitSpec: TypeAlias = _TD64UnitT | tuple[_TD64UnitT, SupportsIndex]
1034+
10131035
### Protocols (for internal use only)
10141036

10151037
@type_check_only
@@ -4010,68 +4032,238 @@ csingle: TypeAlias = complexfloating[_NBitSingle, _NBitSingle]
40104032
cdouble: TypeAlias = complexfloating[_NBitDouble, _NBitDouble]
40114033
clongdouble: TypeAlias = complexfloating[_NBitLongDouble, _NBitLongDouble]
40124034

4013-
class timedelta64(_IntegralMixin, generic[dt.timedelta | int | None]):
4014-
def __init__(
4015-
self,
4016-
value: timedelta64 | dt.timedelta | int | _CharLike_co | None = None,
4017-
format: _CharLike_co | tuple[_CharLike_co, _IntLike_co] = ...,
4018-
/,
4019-
) -> None: ...
4020-
4035+
class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co]):
40214036
@property
40224037
def itemsize(self) -> L[8]: ...
40234038
@property
40244039
def nbytes(self) -> L[8]: ...
40254040

4041+
@overload
4042+
def __init__(self, value: _TD64ItemT_co | timedelta64[_TD64ItemT_co], /) -> None: ...
4043+
@overload
4044+
def __init__(self: timedelta64[L[0]], /) -> None: ...
4045+
@overload
4046+
def __init__(self: timedelta64[None], value: _NaTValue | None, format: _TimeUnitSpec, /) -> None: ...
4047+
@overload
4048+
def __init__(self: timedelta64[int], value: dt.timedelta, format: _TimeUnitSpec[_IntTimeUnit], /) -> None: ...
4049+
@overload
4050+
def __init__(self: timedelta64[int], value: int, format: _TimeUnitSpec[_IntTD64Unit] = ..., /) -> None: ...
4051+
@overload
4052+
def __init__(
4053+
self: timedelta64[dt.timedelta],
4054+
value: dt.timedelta | int,
4055+
format: _TimeUnitSpec[_NativeTD64Unit] = ...,
4056+
/,
4057+
) -> None: ...
4058+
@overload
4059+
def __init__(self, value: int | bytes | str | dt.timedelta | None, format: _TimeUnitSpec = ..., /) -> None: ...
4060+
40264061
# NOTE: Only a limited number of units support conversion
40274062
# to builtin scalar types: `Y`, `M`, `ns`, `ps`, `fs`, `as`
4028-
def __neg__(self) -> Self: ...
4029-
def __pos__(self) -> Self: ...
4030-
def __abs__(self) -> Self: ...
4063+
def __int__(self: timedelta64[int], /) -> int: ...
4064+
def __float__(self: timedelta64[int], /) -> float: ...
40314065

4032-
def __add__(self, other: _TD64Like_co, /) -> timedelta64: ...
4033-
def __radd__(self, other: _TD64Like_co, /) -> timedelta64: ...
4034-
def __sub__(self, other: _TD64Like_co, /) -> timedelta64: ...
4035-
def __rsub__(self, other: _TD64Like_co, /) -> timedelta64: ...
4036-
def __mul__(self, other: _FloatLike_co, /) -> timedelta64: ...
4037-
def __rmul__(self, other: _FloatLike_co, /) -> timedelta64: ...
4038-
__truediv__: _TD64Div[float64]
4039-
def __rtruediv__(self, other: timedelta64, /) -> float64: ...
4040-
__floordiv__: _TD64Div[int64]
4041-
def __rfloordiv__(self, other: timedelta64, /) -> int64: ...
4042-
def __mod__(self, other: timedelta64, /) -> timedelta64: ...
4043-
def __rmod__(self, other: timedelta64, /) -> timedelta64: ...
4044-
def __divmod__(self, other: timedelta64, /) -> tuple[int64, timedelta64]: ...
4045-
def __rdivmod__(self, other: timedelta64, /) -> tuple[int64, timedelta64]: ...
4066+
def __neg__(self, /) -> Self: ...
4067+
def __pos__(self, /) -> Self: ...
4068+
def __abs__(self, /) -> Self: ...
4069+
4070+
@overload
4071+
def __add__(self: timedelta64[None], x: _TD64Like_co, /) -> timedelta64[None]: ...
4072+
@overload
4073+
def __add__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int]: ...
4074+
@overload
4075+
def __add__(self: timedelta64[int], x: timedelta64, /) -> timedelta64[int | None]: ...
4076+
@overload
4077+
def __add__(self: timedelta64[dt.timedelta], x: _AnyDateOrTime, /) -> _AnyDateOrTime: ...
4078+
@overload
4079+
def __add__(self: timedelta64[_AnyTD64Item], x: timedelta64[_AnyTD64Item] | _IntLike_co, /) -> timedelta64[_AnyTD64Item]: ...
4080+
@overload
4081+
def __add__(self, x: timedelta64[None], /) -> timedelta64[None]: ...
4082+
__radd__ = __add__
4083+
4084+
@overload
4085+
def __mul__(self: timedelta64[_AnyTD64Item], x: int | np.integer[Any] | np.bool, /) -> timedelta64[_AnyTD64Item]: ...
4086+
@overload
4087+
def __mul__(self: timedelta64[_AnyTD64Item], x: float | np.floating[Any], /) -> timedelta64[_AnyTD64Item | None]: ...
4088+
@overload
4089+
def __mul__(self, x: float | np.floating[Any] | np.integer[Any] | np.bool, /) -> timedelta64: ...
4090+
__rmul__ = __mul__
4091+
4092+
@overload
4093+
def __mod__(self: timedelta64[None], x: timedelta64, /) -> timedelta64[None]: ...
4094+
@overload
4095+
def __mod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> dt.timedelta: ...
4096+
@overload
4097+
def __mod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item]: ...
4098+
@overload
4099+
def __mod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int]: ...
4100+
@overload
4101+
def __mod__(self, x: timedelta64[None], /) -> timedelta64[None]: ...
4102+
@overload
4103+
def __mod__(self, x: timedelta64[int], /) -> timedelta64[int]: ...
4104+
@overload
4105+
def __mod__(self, x: timedelta64, /) -> timedelta64: ...
4106+
__rmod__ = __mod__ # at runtime the outcomes differ, but the type signatures are the same
4107+
4108+
@overload
4109+
def __divmod__(self: timedelta64[None], x: timedelta64, /) -> tuple[int64, timedelta64[None]]: ...
4110+
@overload
4111+
def __divmod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> tuple[int, dt.timedelta]: ...
4112+
@overload
4113+
def __divmod__(self, x: timedelta64, /) -> tuple[int64, timedelta64]: ...
4114+
__rdivmod__ = __divmod__
4115+
4116+
@overload
4117+
def __sub__(self: timedelta64[None], b: _TD64Like_co, /) -> timedelta64[None]: ...
4118+
@overload
4119+
def __sub__(self: timedelta64[int], b: timedelta64[int | dt.timedelta], /) -> timedelta64[int]: ...
4120+
@overload
4121+
def __sub__(self: timedelta64[int], b: timedelta64, /) -> timedelta64[int | None]: ...
4122+
@overload
4123+
def __sub__(self: timedelta64[dt.timedelta], b: dt.timedelta, /) -> dt.timedelta: ...
4124+
@overload
4125+
def __sub__(self: timedelta64[_AnyTD64Item], b: timedelta64[_AnyTD64Item] | _IntLike_co, /) -> timedelta64[_AnyTD64Item]: ...
4126+
@overload
4127+
def __sub__(self, b: timedelta64[None], /) -> timedelta64[None]: ...
4128+
4129+
@overload
4130+
def __rsub__(self: timedelta64[None], a: _TD64Like_co, /) -> timedelta64[None]: ...
4131+
@overload
4132+
def __rsub__(self: timedelta64[dt.timedelta], a: _AnyDateOrTime, /) -> _AnyDateOrTime: ...
4133+
@overload
4134+
def __rsub__(self: timedelta64[dt.timedelta], a: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item]: ...
4135+
@overload
4136+
def __rsub__(self: timedelta64[_AnyTD64Item], a: timedelta64[_AnyTD64Item] | _IntLike_co, /) -> timedelta64[_AnyTD64Item]: ...
4137+
@overload
4138+
def __rsub__(self, a: timedelta64[None], /) -> timedelta64[None]: ...
4139+
@overload
4140+
def __rsub__(self, a: datetime64[None], /) -> datetime64[None]: ...
4141+
4142+
@overload
4143+
def __truediv__(self: timedelta64[dt.timedelta], b: dt.timedelta, /) -> float: ...
4144+
@overload
4145+
def __truediv__(self, b: timedelta64, /) -> float64: ...
4146+
@overload
4147+
def __truediv__(self: timedelta64[_AnyTD64Item], b: int | integer, /) -> timedelta64[_AnyTD64Item]: ...
4148+
@overload
4149+
def __truediv__(self: timedelta64[_AnyTD64Item], b: float | floating, /) -> timedelta64[_AnyTD64Item | None]: ...
4150+
@overload
4151+
def __truediv__(self, b: float | floating | integer, /) -> timedelta64: ...
4152+
@overload
4153+
def __rtruediv__(self: timedelta64[dt.timedelta], a: dt.timedelta, /) -> float: ...
4154+
@overload
4155+
def __rtruediv__(self, a: timedelta64, /) -> float64: ...
4156+
4157+
@overload
4158+
def __floordiv__(self: timedelta64[dt.timedelta], b: dt.timedelta, /) -> int: ...
4159+
@overload
4160+
def __floordiv__(self, b: timedelta64, /) -> int64: ...
4161+
@overload
4162+
def __floordiv__(self: timedelta64[_AnyTD64Item], b: int | integer, /) -> timedelta64[_AnyTD64Item]: ...
4163+
@overload
4164+
def __floordiv__(self: timedelta64[_AnyTD64Item], b: float | floating, /) -> timedelta64[_AnyTD64Item | None]: ...
4165+
@overload
4166+
def __rfloordiv__(self: timedelta64[dt.timedelta], a: dt.timedelta, /) -> int: ...
4167+
@overload
4168+
def __rfloordiv__(self, a: timedelta64, /) -> int64: ...
40464169

40474170
__lt__: _ComparisonOpLT[_TD64Like_co, _ArrayLikeTD64_co]
40484171
__le__: _ComparisonOpLE[_TD64Like_co, _ArrayLikeTD64_co]
40494172
__gt__: _ComparisonOpGT[_TD64Like_co, _ArrayLikeTD64_co]
40504173
__ge__: _ComparisonOpGE[_TD64Like_co, _ArrayLikeTD64_co]
40514174

4052-
class datetime64(_RealMixin, generic[dt.date | int | None]):
4053-
@overload
4054-
def __init__(self, value: int, format: _CharLike_co | tuple[_CharLike_co, _IntLike_co], /) -> None: ...
4055-
@overload
4056-
def __init__(
4057-
self,
4058-
value: datetime64 | _CharLike_co | _HasDateAttributes | None = None,
4059-
format: _CharLike_co | tuple[_CharLike_co, _IntLike_co] = ...,
4060-
/,
4061-
) -> None: ...
4062-
4175+
class datetime64(_RealMixin, generic[_DT64ItemT_co], Generic[_DT64ItemT_co]):
40634176
@property
40644177
def itemsize(self) -> L[8]: ...
40654178
@property
40664179
def nbytes(self) -> L[8]: ...
40674180

4068-
def __add__(self, other: _TD64Like_co, /) -> datetime64: ...
4069-
def __radd__(self, other: _TD64Like_co, /) -> datetime64: ...
40704181
@overload
4071-
def __sub__(self, other: datetime64, /) -> timedelta64: ...
4182+
def __init__(self, value: datetime64[_DT64ItemT_co], /) -> None: ...
4183+
@overload
4184+
def __init__(self: datetime64[_AnyDT64Arg], value: _AnyDT64Arg, /) -> None: ...
4185+
@overload
4186+
def __init__(self: datetime64[None], value: _NaTValue | None = ..., format: _TimeUnitSpec = ..., /) -> None: ...
4187+
@overload
4188+
def __init__(self: datetime64[dt.datetime], value: _DT64Now, format: _TimeUnitSpec[_NativeTimeUnit] = ..., /) -> None: ...
4189+
@overload
4190+
def __init__(self: datetime64[dt.date], value: _DT64Date, format: _TimeUnitSpec[_DateUnit] = ..., /) -> None: ...
4191+
@overload
4192+
def __init__(self: datetime64[int], value: int | bytes | str | dt.date, format: _TimeUnitSpec[_IntTimeUnit], /) -> None: ...
4193+
@overload
4194+
def __init__(self: datetime64[dt.datetime], value: int | bytes | str, format: _TimeUnitSpec[_NativeTimeUnit], /) -> None: ...
4195+
@overload
4196+
def __init__(self: datetime64[dt.date], value: int | bytes | str, format: _TimeUnitSpec[_DateUnit], /) -> None: ...
4197+
@overload
4198+
def __init__(self, value: bytes | str | None, format: _TimeUnitSpec = ..., /) -> None: ...
4199+
4200+
@overload
4201+
def __add__(self: datetime64[_AnyDT64Item], x: int | integer[Any] | np.bool, /) -> datetime64[_AnyDT64Item]: ...
4202+
@overload
4203+
def __add__(self: datetime64[None], x: _TD64Like_co, /) -> datetime64[None]: ...
4204+
@overload
4205+
def __add__(self: datetime64[int], x: timedelta64[int | dt.timedelta], /) -> datetime64[int]: ...
4206+
@overload
4207+
def __add__(self: datetime64[dt.datetime], x: timedelta64[dt.timedelta], /) -> datetime64[dt.datetime]: ...
4208+
@overload
4209+
def __add__(self: datetime64[dt.date], x: timedelta64[dt.timedelta], /) -> datetime64[dt.date]: ...
4210+
@overload
4211+
def __add__(self: datetime64[dt.date], x: timedelta64[int], /) -> datetime64[int]: ...
4212+
@overload
4213+
def __add__(self, x: datetime64[None], /) -> datetime64[None]: ...
4214+
@overload
4215+
def __add__(self, x: _TD64Like_co, /) -> datetime64: ...
4216+
__radd__ = __add__
4217+
4218+
@overload
4219+
def __sub__(self: datetime64[_AnyDT64Item], x: int | integer[Any] | np.bool, /) -> datetime64[_AnyDT64Item]: ...
4220+
@overload
4221+
def __sub__(self: datetime64[_AnyDate], x: _AnyDate, /) -> dt.timedelta: ...
4222+
@overload
4223+
def __sub__(self: datetime64[None], x: timedelta64, /) -> datetime64[None]: ...
4224+
@overload
4225+
def __sub__(self: datetime64[None], x: datetime64, /) -> timedelta64[None]: ...
4226+
@overload
4227+
def __sub__(self: datetime64[int], x: timedelta64, /) -> datetime64[int]: ...
4228+
@overload
4229+
def __sub__(self: datetime64[int], x: datetime64, /) -> timedelta64[int]: ...
4230+
@overload
4231+
def __sub__(self: datetime64[dt.datetime], x: timedelta64[int], /) -> datetime64[int]: ...
4232+
@overload
4233+
def __sub__(self: datetime64[dt.datetime], x: timedelta64[dt.timedelta], /) -> datetime64[dt.datetime]: ...
4234+
@overload
4235+
def __sub__(self: datetime64[dt.datetime], x: datetime64[int], /) -> timedelta64[int]: ...
4236+
@overload
4237+
def __sub__(self: datetime64[dt.date], x: timedelta64[int], /) -> datetime64[dt.date | int]: ...
4238+
@overload
4239+
def __sub__(self: datetime64[dt.date], x: timedelta64[dt.timedelta], /) -> datetime64[dt.date]: ...
4240+
@overload
4241+
def __sub__(self: datetime64[dt.date], x: datetime64[dt.date], /) -> timedelta64[dt.timedelta]: ...
4242+
@overload
4243+
def __sub__(self, x: timedelta64[None], /) -> datetime64[None]: ...
4244+
@overload
4245+
def __sub__(self, x: datetime64[None], /) -> timedelta64[None]: ...
4246+
@overload
4247+
def __sub__(self, x: _TD64Like_co, /) -> datetime64: ...
4248+
@overload
4249+
def __sub__(self, x: datetime64, /) -> timedelta64: ...
4250+
4251+
@overload
4252+
def __rsub__(self: datetime64[_AnyDT64Item], x: int | integer[Any] | np.bool, /) -> datetime64[_AnyDT64Item]: ...
4253+
@overload
4254+
def __rsub__(self: datetime64[_AnyDate], x: _AnyDate, /) -> dt.timedelta: ...
4255+
@overload
4256+
def __rsub__(self: datetime64[None], x: datetime64, /) -> timedelta64[None]: ...
4257+
@overload
4258+
def __rsub__(self: datetime64[int], x: datetime64, /) -> timedelta64[int]: ...
4259+
@overload
4260+
def __rsub__(self: datetime64[dt.datetime], x: datetime64[int], /) -> timedelta64[int]: ...
4261+
@overload
4262+
def __rsub__(self: datetime64[dt.datetime], x: datetime64[dt.date], /) -> timedelta64[dt.timedelta]: ...
4263+
@overload
4264+
def __rsub__(self, x: datetime64[None], /) -> timedelta64[None]: ...
40724265
@overload
4073-
def __sub__(self, other: _TD64Like_co, /) -> datetime64: ...
4074-
def __rsub__(self, other: datetime64, /) -> timedelta64: ...
4266+
def __rsub__(self, x: datetime64, /) -> timedelta64: ...
40754267

40764268
__lt__: _ComparisonOpLT[datetime64, _ArrayLikeDT64_co]
40774269
__le__: _ComparisonOpLE[datetime64, _ArrayLikeDT64_co]

‎numpy/_typing/_callable.pyi

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@ from typing import (
2222
import numpy as np
2323
from numpy import (
2424
generic,
25-
timedelta64,
2625
number,
2726
integer,
2827
unsignedinteger,
@@ -38,7 +37,6 @@ from ._nbit import _NBitInt
3837
from ._scalars import (
3938
_BoolLike_co,
4039
_IntLike_co,
41-
_FloatLike_co,
4240
_NumberLike_co,
4341
)
4442
from . import NBitBase
@@ -132,15 +130,6 @@ class _BoolDivMod(Protocol):
132130
@overload
133131
def __call__(self, other: _FloatType, /) -> _2Tuple[_FloatType]: ...
134132

135-
@type_check_only
136-
class _TD64Div(Protocol[_NumberType_co]):
137-
@overload
138-
def __call__(self, other: timedelta64, /) -> _NumberType_co: ...
139-
@overload
140-
def __call__(self, other: _BoolLike_co, /) -> NoReturn: ...
141-
@overload
142-
def __call__(self, other: _FloatLike_co, /) -> timedelta64: ...
143-
144133
@type_check_only
145134
class _IntTrueDiv(Protocol[_NBit1]):
146135
@overload

0 commit comments

Comments
 (0)