1
1
from __future__ import annotations
2
2
3
- import datetime
4
3
import numbers
5
4
from typing import (
6
5
TYPE_CHECKING ,
10
9
import numpy as np
11
10
12
11
from pandas ._libs import (
13
- Timedelta ,
14
12
lib ,
15
13
missing as libmissing ,
16
14
)
23
21
24
22
from pandas .core .dtypes .common import (
25
23
is_bool_dtype ,
26
- is_float ,
27
24
is_float_dtype ,
28
- is_integer ,
29
25
is_integer_dtype ,
30
26
is_list_like ,
31
27
is_object_dtype ,
32
28
is_string_dtype ,
33
29
pandas_dtype ,
34
30
)
35
31
32
+ from pandas .core import ops
33
+ from pandas .core .arrays .base import ExtensionArray
36
34
from pandas .core .arrays .masked import (
37
35
BaseMaskedArray ,
38
36
BaseMaskedDtype ,
39
37
)
38
+ from pandas .core .construction import ensure_wrapped_if_datetimelike
40
39
41
40
if TYPE_CHECKING :
42
41
import pyarrow
@@ -219,34 +218,40 @@ def _arith_method(self, other, op):
219
218
op_name = op .__name__
220
219
omask = None
221
220
222
- if getattr (other , "ndim" , 0 ) > 1 :
223
- raise NotImplementedError ("can only perform ops with 1-d structures" )
224
-
225
- if isinstance (other , NumericArray ):
221
+ if isinstance (other , BaseMaskedArray ):
226
222
other , omask = other ._data , other ._mask
227
223
228
224
elif is_list_like (other ):
229
- other = np .asarray (other )
225
+ if not isinstance (other , ExtensionArray ):
226
+ other = np .asarray (other )
230
227
if other .ndim > 1 :
231
228
raise NotImplementedError ("can only perform ops with 1-d structures" )
232
- if len (self ) != len (other ):
233
- raise ValueError ("Lengths must match" )
234
- if not (is_float_dtype (other ) or is_integer_dtype (other )):
235
- raise TypeError ("can only perform ops with numeric values" )
236
229
237
- elif isinstance (other , (datetime .timedelta , np .timedelta64 )):
238
- other = Timedelta (other )
230
+ # We wrap the non-masked arithmetic logic used for numpy dtypes
231
+ # in Series/Index arithmetic ops.
232
+ other = ops .maybe_prepare_scalar_for_op (other , (len (self ),))
233
+ pd_op = ops .get_array_op (op )
234
+ other = ensure_wrapped_if_datetimelike (other )
239
235
240
- else :
241
- if not (is_float (other ) or is_integer (other ) or other is libmissing .NA ):
242
- raise TypeError ("can only perform ops with numeric values" )
236
+ mask = self ._propagate_mask (omask , other )
243
237
244
- if omask is None :
245
- mask = self ._mask .copy ()
246
- if other is libmissing .NA :
247
- mask |= True
238
+ if other is libmissing .NA :
239
+ result = np .ones_like (self ._data )
240
+ if "truediv" in op_name and self .dtype .kind != "f" :
241
+ # The actual data here doesn't matter since the mask
242
+ # will be all-True, but since this is division, we want
243
+ # to end up with floating dtype.
244
+ result = result .astype (np .float64 )
248
245
else :
249
- mask = self ._mask | omask
246
+ # Make sure we do this before the "pow" mask checks
247
+ # to get an expected exception message on shape mismatch.
248
+ if self .dtype .kind in ["i" , "u" ] and op_name in ["floordiv" , "mod" ]:
249
+ # ATM we don't match the behavior of non-masked types with
250
+ # respect to floordiv-by-zero
251
+ pd_op = op
252
+
253
+ with np .errstate (all = "ignore" ):
254
+ result = pd_op (self ._data , other )
250
255
251
256
if op_name == "pow" :
252
257
# 1 ** x is 1.
@@ -266,25 +271,6 @@ def _arith_method(self, other, op):
266
271
# x ** 0 is 1.
267
272
mask = np .where ((self ._data == 0 ) & ~ self ._mask , False , mask )
268
273
269
- if other is libmissing .NA :
270
- result = np .ones_like (self ._data )
271
- if "truediv" in op_name and self .dtype .kind != "f" :
272
- # The actual data here doesn't matter since the mask
273
- # will be all-True, but since this is division, we want
274
- # to end up with floating dtype.
275
- result = result .astype (np .float64 )
276
- else :
277
- with np .errstate (all = "ignore" ):
278
- result = op (self ._data , other )
279
-
280
- # divmod returns a tuple
281
- if op_name == "divmod" :
282
- div , mod = result
283
- return (
284
- self ._maybe_mask_result (div , mask , other , "floordiv" ),
285
- self ._maybe_mask_result (mod , mask , other , "mod" ),
286
- )
287
-
288
274
return self ._maybe_mask_result (result , mask , other , op_name )
289
275
290
276
_HANDLED_TYPES = (np .ndarray , numbers .Number )
0 commit comments