72
72
from pandas .util ._exceptions import find_stack_level
73
73
74
74
from pandas .core .dtypes .common import (
75
+ DT64NS_DTYPE ,
75
76
is_all_strings ,
76
77
is_categorical_dtype ,
77
78
is_datetime64_any_dtype ,
89
90
is_unsigned_integer_dtype ,
90
91
pandas_dtype ,
91
92
)
92
- from pandas .core .dtypes .dtypes import ExtensionDtype
93
+ from pandas .core .dtypes .dtypes import (
94
+ DatetimeTZDtype ,
95
+ ExtensionDtype ,
96
+ )
93
97
from pandas .core .dtypes .missing import (
94
98
is_valid_na_for_dtype ,
95
99
isna ,
113
117
import pandas .core .common as com
114
118
from pandas .core .construction import (
115
119
array as pd_array ,
120
+ ensure_wrapped_if_datetimelike ,
116
121
extract_array ,
117
122
)
118
123
from pandas .core .indexers import (
@@ -1082,26 +1087,123 @@ def _cmp_method(self, other, op):
1082
1087
__divmod__ = make_invalid_op ("__divmod__" )
1083
1088
__rdivmod__ = make_invalid_op ("__rdivmod__" )
1084
1089
1090
+ @final
1085
1091
def _add_datetimelike_scalar (self , other ):
1086
1092
# 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
+ )
1088
1097
1089
- _add_datetime_arraylike = _add_datetimelike_scalar
1098
+ from pandas . core . arrays import DatetimeArray
1090
1099
1091
- def _sub_datetimelike_scalar (self , other ):
1092
- # Overridden by DatetimeArray
1093
1100
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
1095
1142
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" )
1097
1161
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
1098
1180
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
1101
1195
1196
+ @final
1102
1197
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
1105
1207
1106
1208
def _add_offset (self , offset ):
1107
1209
raise AbstractMethodError (self )
@@ -1116,9 +1218,9 @@ def _add_timedeltalike_scalar(self, other):
1116
1218
"""
1117
1219
if isna (other ):
1118
1220
# 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 )
1120
1222
new_values .fill (iNaT )
1121
- return type (self )(new_values , dtype = self .dtype )
1223
+ return type (self ). _simple_new (new_values , dtype = self .dtype )
1122
1224
1123
1225
# PeriodArray overrides, so we only get here with DTA/TDA
1124
1226
# error: "DatetimeLikeArrayMixin" has no attribute "_reso"
@@ -1139,7 +1241,9 @@ def _add_timedeltalike_scalar(self, other):
1139
1241
new_values , dtype = self .dtype , freq = new_freq
1140
1242
)
1141
1243
1142
- def _add_timedelta_arraylike (self , other ):
1244
+ def _add_timedelta_arraylike (
1245
+ self , other : TimedeltaArray | npt .NDArray [np .timedelta64 ]
1246
+ ):
1143
1247
"""
1144
1248
Add a delta of a TimedeltaIndex
1145
1249
@@ -1152,11 +1256,8 @@ def _add_timedelta_arraylike(self, other):
1152
1256
if len (self ) != len (other ):
1153
1257
raise ValueError ("cannot add indices of unequal length" )
1154
1258
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 )
1160
1261
1161
1262
self_i8 = self .asi8
1162
1263
other_i8 = other .asi8
@@ -1200,12 +1301,27 @@ def _sub_nat(self):
1200
1301
result .fill (iNaT )
1201
1302
return result .view ("timedelta64[ns]" )
1202
1303
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
1207
1316
)
1208
1317
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
1209
1325
def _addsub_object_array (self , other : np .ndarray , op ):
1210
1326
"""
1211
1327
Add or subtract array-like of DateOffset objects
0 commit comments