93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
import datetime
|
|
|
|
import numpy as np
|
|
|
|
from pandas._libs import Timedelta, missing as libmissing
|
|
from pandas.errors import AbstractMethodError
|
|
|
|
from pandas.core.dtypes.common import (
|
|
is_float,
|
|
is_float_dtype,
|
|
is_integer,
|
|
is_integer_dtype,
|
|
is_list_like,
|
|
)
|
|
|
|
from .masked import BaseMaskedArray
|
|
|
|
|
|
class NumericArray(BaseMaskedArray):
|
|
"""
|
|
Base class for IntegerArray and FloatingArray.
|
|
"""
|
|
|
|
def _maybe_mask_result(self, result, mask, other, op_name: str):
|
|
raise AbstractMethodError(self)
|
|
|
|
def _arith_method(self, other, op):
|
|
op_name = op.__name__
|
|
omask = None
|
|
|
|
if getattr(other, "ndim", 0) > 1:
|
|
raise NotImplementedError("can only perform ops with 1-d structures")
|
|
|
|
if isinstance(other, NumericArray):
|
|
other, omask = other._data, other._mask
|
|
|
|
elif is_list_like(other):
|
|
other = np.asarray(other)
|
|
if other.ndim > 1:
|
|
raise NotImplementedError("can only perform ops with 1-d structures")
|
|
if len(self) != len(other):
|
|
raise ValueError("Lengths must match")
|
|
if not (is_float_dtype(other) or is_integer_dtype(other)):
|
|
raise TypeError("can only perform ops with numeric values")
|
|
|
|
elif isinstance(other, (datetime.timedelta, np.timedelta64)):
|
|
other = Timedelta(other)
|
|
|
|
else:
|
|
if not (is_float(other) or is_integer(other) or other is libmissing.NA):
|
|
raise TypeError("can only perform ops with numeric values")
|
|
|
|
if omask is None:
|
|
mask = self._mask.copy()
|
|
if other is libmissing.NA:
|
|
mask |= True
|
|
else:
|
|
mask = self._mask | omask
|
|
|
|
if op_name == "pow":
|
|
# 1 ** x is 1.
|
|
mask = np.where((self._data == 1) & ~self._mask, False, mask)
|
|
# x ** 0 is 1.
|
|
if omask is not None:
|
|
mask = np.where((other == 0) & ~omask, False, mask)
|
|
elif other is not libmissing.NA:
|
|
mask = np.where(other == 0, False, mask)
|
|
|
|
elif op_name == "rpow":
|
|
# 1 ** x is 1.
|
|
if omask is not None:
|
|
mask = np.where((other == 1) & ~omask, False, mask)
|
|
elif other is not libmissing.NA:
|
|
mask = np.where(other == 1, False, mask)
|
|
# x ** 0 is 1.
|
|
mask = np.where((self._data == 0) & ~self._mask, False, mask)
|
|
|
|
if other is libmissing.NA:
|
|
result = np.ones_like(self._data)
|
|
else:
|
|
with np.errstate(all="ignore"):
|
|
result = op(self._data, other)
|
|
|
|
# divmod returns a tuple
|
|
if op_name == "divmod":
|
|
div, mod = result
|
|
return (
|
|
self._maybe_mask_result(div, mask, other, "floordiv"),
|
|
self._maybe_mask_result(mod, mask, other, "mod"),
|
|
)
|
|
|
|
return self._maybe_mask_result(result, mask, other, op_name)
|