Skip to content

Commit 3350f95

Browse files
authored
REF: share DTA/TDA/PA arithmetic methods (#47205)
* REF: share DTA/TDA/PA arithmetic methods * troubleshoot npdev build
1 parent 5f72083 commit 3350f95

File tree

5 files changed

+176
-204
lines changed

5 files changed

+176
-204
lines changed

pandas/_libs/tslibs/period.pyx

+11-13
Original file line numberDiff line numberDiff line change
@@ -1728,10 +1728,12 @@ cdef class _Period(PeriodMixin):
17281728
elif util.is_integer_object(other):
17291729
ordinal = self.ordinal + other * self.freq.n
17301730
return Period(ordinal=ordinal, freq=self.freq)
1731-
elif (PyDateTime_Check(other) or
1732-
is_period_object(other) or util.is_datetime64_object(other)):
1731+
1732+
elif is_period_object(other):
17331733
# can't add datetime-like
1734-
# GH#17983
1734+
# GH#17983; can't just return NotImplemented bc we get a RecursionError
1735+
# when called via np.add.reduce see TestNumpyReductions.test_add
1736+
# in npdev build
17351737
sname = type(self).__name__
17361738
oname = type(other).__name__
17371739
raise TypeError(f"unsupported operand type(s) for +: '{sname}' "
@@ -1750,16 +1752,12 @@ cdef class _Period(PeriodMixin):
17501752
return NaT
17511753
return NotImplemented
17521754

1753-
elif is_any_td_scalar(other):
1754-
neg_other = -other
1755-
return self + neg_other
1756-
elif is_offset_object(other):
1757-
# Non-Tick DateOffset
1758-
neg_other = -other
1759-
return self + neg_other
1760-
elif util.is_integer_object(other):
1761-
ordinal = self.ordinal - other * self.freq.n
1762-
return Period(ordinal=ordinal, freq=self.freq)
1755+
elif (
1756+
is_any_td_scalar(other)
1757+
or is_offset_object(other)
1758+
or util.is_integer_object(other)
1759+
):
1760+
return self + (-other)
17631761
elif is_period_object(other):
17641762
self._require_matching_freq(other)
17651763
# GH 23915 - mul by base freq since __add__ is agnostic of n

pandas/core/arrays/datetimelike.py

+139-23
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@
7272
from pandas.util._exceptions import find_stack_level
7373

7474
from pandas.core.dtypes.common import (
75+
DT64NS_DTYPE,
7576
is_all_strings,
7677
is_categorical_dtype,
7778
is_datetime64_any_dtype,
@@ -89,7 +90,10 @@
8990
is_unsigned_integer_dtype,
9091
pandas_dtype,
9192
)
92-
from pandas.core.dtypes.dtypes import ExtensionDtype
93+
from pandas.core.dtypes.dtypes import (
94+
DatetimeTZDtype,
95+
ExtensionDtype,
96+
)
9397
from pandas.core.dtypes.missing import (
9498
is_valid_na_for_dtype,
9599
isna,
@@ -113,6 +117,7 @@
113117
import pandas.core.common as com
114118
from pandas.core.construction import (
115119
array as pd_array,
120+
ensure_wrapped_if_datetimelike,
116121
extract_array,
117122
)
118123
from pandas.core.indexers import (
@@ -1082,26 +1087,123 @@ def _cmp_method(self, other, op):
10821087
__divmod__ = make_invalid_op("__divmod__")
10831088
__rdivmod__ = make_invalid_op("__rdivmod__")
10841089

1090+
@final
10851091
def _add_datetimelike_scalar(self, other):
10861092
# Overridden by TimedeltaArray
1087-
raise TypeError(f"cannot add {type(self).__name__} and {type(other).__name__}")
1093+
if not is_timedelta64_dtype(self.dtype):
1094+
raise TypeError(
1095+
f"cannot add {type(self).__name__} and {type(other).__name__}"
1096+
)
10881097

1089-
_add_datetime_arraylike = _add_datetimelike_scalar
1098+
from pandas.core.arrays import DatetimeArray
10901099

1091-
def _sub_datetimelike_scalar(self, other):
1092-
# Overridden by DatetimeArray
10931100
assert other is not NaT
1094-
raise TypeError(f"cannot subtract a datelike from a {type(self).__name__}")
1101+
other = Timestamp(other)
1102+
if other is NaT:
1103+
# In this case we specifically interpret NaT as a datetime, not
1104+
# the timedelta interpretation we would get by returning self + NaT
1105+
result = self.asi8.view("m8[ms]") + NaT.to_datetime64()
1106+
return DatetimeArray(result)
1107+
1108+
i8 = self.asi8
1109+
result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan)
1110+
result = self._maybe_mask_results(result)
1111+
dtype = DatetimeTZDtype(tz=other.tz) if other.tz else DT64NS_DTYPE
1112+
return DatetimeArray(result, dtype=dtype, freq=self.freq)
1113+
1114+
@final
1115+
def _add_datetime_arraylike(self, other):
1116+
if not is_timedelta64_dtype(self.dtype):
1117+
raise TypeError(
1118+
f"cannot add {type(self).__name__} and {type(other).__name__}"
1119+
)
1120+
1121+
# At this point we have already checked that other.dtype is datetime64
1122+
other = ensure_wrapped_if_datetimelike(other)
1123+
# defer to DatetimeArray.__add__
1124+
return other + self
1125+
1126+
@final
1127+
def _sub_datetimelike_scalar(self, other: datetime | np.datetime64):
1128+
if self.dtype.kind != "M":
1129+
raise TypeError(f"cannot subtract a datelike from a {type(self).__name__}")
1130+
1131+
self = cast("DatetimeArray", self)
1132+
# subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
1133+
1134+
# error: Non-overlapping identity check (left operand type: "Union[datetime,
1135+
# datetime64]", right operand type: "NaTType") [comparison-overlap]
1136+
assert other is not NaT # type: ignore[comparison-overlap]
1137+
other = Timestamp(other)
1138+
# error: Non-overlapping identity check (left operand type: "Timestamp",
1139+
# right operand type: "NaTType")
1140+
if other is NaT: # type: ignore[comparison-overlap]
1141+
return self - NaT
10951142

1096-
_sub_datetime_arraylike = _sub_datetimelike_scalar
1143+
try:
1144+
self._assert_tzawareness_compat(other)
1145+
except TypeError as err:
1146+
new_message = str(err).replace("compare", "subtract")
1147+
raise type(err)(new_message) from err
1148+
1149+
i8 = self.asi8
1150+
result = checked_add_with_arr(i8, -other.value, arr_mask=self._isnan)
1151+
result = self._maybe_mask_results(result)
1152+
return result.view("timedelta64[ns]")
1153+
1154+
@final
1155+
def _sub_datetime_arraylike(self, other):
1156+
if self.dtype.kind != "M":
1157+
raise TypeError(f"cannot subtract a datelike from a {type(self).__name__}")
1158+
1159+
if len(self) != len(other):
1160+
raise ValueError("cannot add indices of unequal length")
10971161

1162+
self = cast("DatetimeArray", self)
1163+
other = ensure_wrapped_if_datetimelike(other)
1164+
1165+
try:
1166+
self._assert_tzawareness_compat(other)
1167+
except TypeError as err:
1168+
new_message = str(err).replace("compare", "subtract")
1169+
raise type(err)(new_message) from err
1170+
1171+
self_i8 = self.asi8
1172+
other_i8 = other.asi8
1173+
arr_mask = self._isnan | other._isnan
1174+
new_values = checked_add_with_arr(self_i8, -other_i8, arr_mask=arr_mask)
1175+
if self._hasna or other._hasna:
1176+
np.putmask(new_values, arr_mask, iNaT)
1177+
return new_values.view("timedelta64[ns]")
1178+
1179+
@final
10981180
def _sub_period(self, other: Period):
1099-
# Overridden by PeriodArray
1100-
raise TypeError(f"cannot subtract Period from a {type(self).__name__}")
1181+
if not is_period_dtype(self.dtype):
1182+
raise TypeError(f"cannot subtract Period from a {type(self).__name__}")
1183+
1184+
# If the operation is well-defined, we return an object-dtype ndarray
1185+
# of DateOffsets. Null entries are filled with pd.NaT
1186+
self._check_compatible_with(other)
1187+
asi8 = self.asi8
1188+
new_data = asi8 - other.ordinal
1189+
new_data = np.array([self.freq.base * x for x in new_data])
1190+
1191+
if self._hasna:
1192+
new_data[self._isnan] = NaT
1193+
1194+
return new_data
11011195

1196+
@final
11021197
def _add_period(self, other: Period):
1103-
# Overridden by TimedeltaArray
1104-
raise TypeError(f"cannot add Period to a {type(self).__name__}")
1198+
if not is_timedelta64_dtype(self.dtype):
1199+
raise TypeError(f"cannot add Period to a {type(self).__name__}")
1200+
1201+
# We will wrap in a PeriodArray and defer to the reversed operation
1202+
from pandas.core.arrays.period import PeriodArray
1203+
1204+
i8vals = np.broadcast_to(other.ordinal, self.shape)
1205+
parr = PeriodArray(i8vals, freq=other.freq)
1206+
return parr + self
11051207

11061208
def _add_offset(self, offset):
11071209
raise AbstractMethodError(self)
@@ -1116,9 +1218,9 @@ def _add_timedeltalike_scalar(self, other):
11161218
"""
11171219
if isna(other):
11181220
# i.e np.timedelta64("NaT"), not recognized by delta_to_nanoseconds
1119-
new_values = np.empty(self.shape, dtype="i8")
1221+
new_values = np.empty(self.shape, dtype="i8").view(self._ndarray.dtype)
11201222
new_values.fill(iNaT)
1121-
return type(self)(new_values, dtype=self.dtype)
1223+
return type(self)._simple_new(new_values, dtype=self.dtype)
11221224

11231225
# PeriodArray overrides, so we only get here with DTA/TDA
11241226
# error: "DatetimeLikeArrayMixin" has no attribute "_reso"
@@ -1139,7 +1241,9 @@ def _add_timedeltalike_scalar(self, other):
11391241
new_values, dtype=self.dtype, freq=new_freq
11401242
)
11411243

1142-
def _add_timedelta_arraylike(self, other):
1244+
def _add_timedelta_arraylike(
1245+
self, other: TimedeltaArray | npt.NDArray[np.timedelta64]
1246+
):
11431247
"""
11441248
Add a delta of a TimedeltaIndex
11451249
@@ -1152,11 +1256,8 @@ def _add_timedelta_arraylike(self, other):
11521256
if len(self) != len(other):
11531257
raise ValueError("cannot add indices of unequal length")
11541258

1155-
if isinstance(other, np.ndarray):
1156-
# ndarray[timedelta64]; wrap in TimedeltaIndex for op
1157-
from pandas.core.arrays import TimedeltaArray
1158-
1159-
other = TimedeltaArray._from_sequence(other)
1259+
other = ensure_wrapped_if_datetimelike(other)
1260+
other = cast("TimedeltaArray", other)
11601261

11611262
self_i8 = self.asi8
11621263
other_i8 = other.asi8
@@ -1200,12 +1301,27 @@ def _sub_nat(self):
12001301
result.fill(iNaT)
12011302
return result.view("timedelta64[ns]")
12021303

1203-
def _sub_period_array(self, other):
1204-
# Overridden by PeriodArray
1205-
raise TypeError(
1206-
f"cannot subtract {other.dtype}-dtype from {type(self).__name__}"
1304+
@final
1305+
def _sub_period_array(self, other: PeriodArray) -> npt.NDArray[np.object_]:
1306+
if not is_period_dtype(self.dtype):
1307+
raise TypeError(
1308+
f"cannot subtract {other.dtype}-dtype from {type(self).__name__}"
1309+
)
1310+
1311+
self = cast("PeriodArray", self)
1312+
self._require_matching_freq(other)
1313+
1314+
new_values = checked_add_with_arr(
1315+
self.asi8, -other.asi8, arr_mask=self._isnan, b_mask=other._isnan
12071316
)
12081317

1318+
new_values = np.array([self.freq.base * x for x in new_values])
1319+
if self._hasna or other._hasna:
1320+
mask = self._isnan | other._isnan
1321+
new_values[mask] = NaT
1322+
return new_values
1323+
1324+
@final
12091325
def _addsub_object_array(self, other: np.ndarray, op):
12101326
"""
12111327
Add or subtract array-like of DateOffset objects

pandas/core/arrays/datetimes.py

+16-60
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,6 @@
7070
from pandas.core.dtypes.generic import ABCMultiIndex
7171
from pandas.core.dtypes.missing import isna
7272

73-
from pandas.core.algorithms import checked_add_with_arr
7473
from pandas.core.arrays import (
7574
ExtensionArray,
7675
datetimelike as dtl,
@@ -733,43 +732,17 @@ def _assert_tzawareness_compat(self, other) -> None:
733732
# -----------------------------------------------------------------
734733
# Arithmetic Methods
735734

736-
def _sub_datetime_arraylike(self, other):
737-
"""subtract DatetimeArray/Index or ndarray[datetime64]"""
738-
if len(self) != len(other):
739-
raise ValueError("cannot add indices of unequal length")
740-
741-
if isinstance(other, np.ndarray):
742-
assert is_datetime64_dtype(other)
743-
other = type(self)(other)
744-
745-
try:
746-
self._assert_tzawareness_compat(other)
747-
except TypeError as error:
748-
new_message = str(error).replace("compare", "subtract")
749-
raise type(error)(new_message) from error
750-
751-
self_i8 = self.asi8
752-
other_i8 = other.asi8
753-
arr_mask = self._isnan | other._isnan
754-
new_values = checked_add_with_arr(self_i8, -other_i8, arr_mask=arr_mask)
755-
if self._hasna or other._hasna:
756-
np.putmask(new_values, arr_mask, iNaT)
757-
return new_values.view("timedelta64[ns]")
758-
759735
def _add_offset(self, offset) -> DatetimeArray:
760736

761737
assert not isinstance(offset, Tick)
738+
739+
if self.tz is not None:
740+
values = self.tz_localize(None)
741+
else:
742+
values = self
743+
762744
try:
763-
if self.tz is not None:
764-
values = self.tz_localize(None)
765-
else:
766-
values = self
767745
result = offset._apply_array(values).view(values.dtype)
768-
result = DatetimeArray._simple_new(result, dtype=result.dtype)
769-
if self.tz is not None:
770-
# FIXME: tz_localize with non-nano
771-
result = result.tz_localize(self.tz)
772-
773746
except NotImplementedError:
774747
warnings.warn(
775748
"Non-vectorized DateOffset being applied to Series or DatetimeIndex.",
@@ -781,35 +754,18 @@ def _add_offset(self, offset) -> DatetimeArray:
781754
# GH#30336 _from_sequence won't be able to infer self.tz
782755
return result.tz_localize(self.tz)
783756

784-
return result
785-
786-
def _sub_datetimelike_scalar(self, other):
787-
# subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
788-
assert isinstance(other, (datetime, np.datetime64))
789-
# error: Non-overlapping identity check (left operand type: "Union[datetime,
790-
# datetime64]", right operand type: "NaTType") [comparison-overlap]
791-
assert other is not NaT # type: ignore[comparison-overlap]
792-
other = Timestamp(other)
793-
# error: Non-overlapping identity check (left operand type: "Timestamp",
794-
# right operand type: "NaTType")
795-
if other is NaT: # type: ignore[comparison-overlap]
796-
return self - NaT
797-
798-
try:
799-
self._assert_tzawareness_compat(other)
800-
except TypeError as error:
801-
new_message = str(error).replace("compare", "subtract")
802-
raise type(error)(new_message) from error
757+
else:
758+
result = DatetimeArray._simple_new(result, dtype=result.dtype)
759+
if self.tz is not None:
760+
# FIXME: tz_localize with non-nano
761+
result = result.tz_localize(self.tz)
803762

804-
i8 = self.asi8
805-
result = checked_add_with_arr(i8, -other.value, arr_mask=self._isnan)
806-
result = self._maybe_mask_results(result)
807-
return result.view("timedelta64[ns]")
763+
return result
808764

809765
# -----------------------------------------------------------------
810766
# Timezone Conversion and Localization Methods
811767

812-
def _local_timestamps(self) -> np.ndarray:
768+
def _local_timestamps(self) -> npt.NDArray[np.int64]:
813769
"""
814770
Convert to an i8 (unix-like nanosecond timestamp) representation
815771
while keeping the local timezone and not using UTC.
@@ -1238,7 +1194,7 @@ def to_perioddelta(self, freq) -> TimedeltaArray:
12381194
# -----------------------------------------------------------------
12391195
# Properties - Vectorized Timestamp Properties/Methods
12401196

1241-
def month_name(self, locale=None):
1197+
def month_name(self, locale=None) -> npt.NDArray[np.object_]:
12421198
"""
12431199
Return the month names of the :class:`~pandas.Series` or
12441200
:class:`~pandas.DatetimeIndex` with specified locale.
@@ -1283,7 +1239,7 @@ def month_name(self, locale=None):
12831239
result = self._maybe_mask_results(result, fill_value=None)
12841240
return result
12851241

1286-
def day_name(self, locale=None):
1242+
def day_name(self, locale=None) -> npt.NDArray[np.object_]:
12871243
"""
12881244
Return the day names of the :class:`~pandas.Series` or
12891245
:class:`~pandas.DatetimeIndex` with specified locale.
@@ -1949,7 +1905,7 @@ def weekofyear(self):
19491905
""",
19501906
)
19511907

1952-
def to_julian_date(self) -> np.ndarray:
1908+
def to_julian_date(self) -> npt.NDArray[np.float64]:
19531909
"""
19541910
Convert Datetime Array to float64 ndarray of Julian Dates.
19551911
0 Julian date is noon January 1, 4713 BC.

0 commit comments

Comments
 (0)