Skip to content

Commit

Permalink
REF: use OpsMixin in EAs (pandas-dev#37049)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jbrockmendel
jbrockmendel authored and Kevin D Smith committed Nov 2, 2020
1 parent e4f7bc7 commit 9958b5c
Show file tree
Hide file tree
Showing 6 changed files with 233 additions and 315 deletions.
175 changes: 76 additions & 99 deletions pandas/core/arrays/boolean.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@

from pandas._libs import lib, missing as libmissing
from pandas._typing import ArrayLike
from pandas.compat import set_function_name
from pandas.compat.numpy import function as nv

from pandas.core.dtypes.common import (
Expand All @@ -23,7 +22,6 @@
from pandas.core.dtypes.missing import isna

from pandas.core import ops
from pandas.core.arraylike import OpsMixin

from .masked import BaseMaskedArray, BaseMaskedDtype

Expand Down Expand Up @@ -203,7 +201,7 @@ def coerce_to_array(
return values, mask


class BooleanArray(OpsMixin, BaseMaskedArray):
class BooleanArray(BaseMaskedArray):
"""
Array of boolean (True/False) data with missing values.
Expand Down Expand Up @@ -561,48 +559,40 @@ def all(self, skipna: bool = True, **kwargs):
else:
return self.dtype.na_value

@classmethod
def _create_logical_method(cls, op):
@ops.unpack_zerodim_and_defer(op.__name__)
def logical_method(self, other):

assert op.__name__ in {"or_", "ror_", "and_", "rand_", "xor", "rxor"}
other_is_booleanarray = isinstance(other, BooleanArray)
other_is_scalar = lib.is_scalar(other)
mask = None

if other_is_booleanarray:
other, mask = other._data, other._mask
elif is_list_like(other):
other = np.asarray(other, dtype="bool")
if other.ndim > 1:
raise NotImplementedError(
"can only perform ops with 1-d structures"
)
other, mask = coerce_to_array(other, copy=False)
elif isinstance(other, np.bool_):
other = other.item()

if other_is_scalar and not (other is libmissing.NA or lib.is_bool(other)):
raise TypeError(
"'other' should be pandas.NA or a bool. "
f"Got {type(other).__name__} instead."
)

if not other_is_scalar and len(self) != len(other):
raise ValueError("Lengths must match to compare")
def _logical_method(self, other, op):

assert op.__name__ in {"or_", "ror_", "and_", "rand_", "xor", "rxor"}
other_is_booleanarray = isinstance(other, BooleanArray)
other_is_scalar = lib.is_scalar(other)
mask = None

if other_is_booleanarray:
other, mask = other._data, other._mask
elif is_list_like(other):
other = np.asarray(other, dtype="bool")
if other.ndim > 1:
raise NotImplementedError("can only perform ops with 1-d structures")
other, mask = coerce_to_array(other, copy=False)
elif isinstance(other, np.bool_):
other = other.item()

if op.__name__ in {"or_", "ror_"}:
result, mask = ops.kleene_or(self._data, other, self._mask, mask)
elif op.__name__ in {"and_", "rand_"}:
result, mask = ops.kleene_and(self._data, other, self._mask, mask)
elif op.__name__ in {"xor", "rxor"}:
result, mask = ops.kleene_xor(self._data, other, self._mask, mask)
if other_is_scalar and not (other is libmissing.NA or lib.is_bool(other)):
raise TypeError(
"'other' should be pandas.NA or a bool. "
f"Got {type(other).__name__} instead."
)

return BooleanArray(result, mask)
if not other_is_scalar and len(self) != len(other):
raise ValueError("Lengths must match to compare")

name = f"__{op.__name__}__"
return set_function_name(logical_method, name, cls)
if op.__name__ in {"or_", "ror_"}:
result, mask = ops.kleene_or(self._data, other, self._mask, mask)
elif op.__name__ in {"and_", "rand_"}:
result, mask = ops.kleene_and(self._data, other, self._mask, mask)
elif op.__name__ in {"xor", "rxor"}:
result, mask = ops.kleene_xor(self._data, other, self._mask, mask)

return BooleanArray(result, mask)

def _cmp_method(self, other, op):
from pandas.arrays import FloatingArray, IntegerArray
Expand Down Expand Up @@ -643,6 +633,50 @@ def _cmp_method(self, other, op):

return BooleanArray(result, mask, copy=False)

def _arith_method(self, other, op):
mask = None
op_name = op.__name__

if isinstance(other, BooleanArray):
other, mask = 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")

# nans propagate
if mask is None:
mask = self._mask
if other is libmissing.NA:
mask |= True
else:
mask = self._mask | mask

if other is libmissing.NA:
# if other is NA, the result will be all NA and we can't run the
# actual op, so we need to choose the resulting dtype manually
if op_name in {"floordiv", "rfloordiv", "mod", "rmod", "pow", "rpow"}:
dtype = "int8"
else:
dtype = "bool"
result = np.zeros(len(self._data), dtype=dtype)
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)

def _reduce(self, name: str, skipna: bool = True, **kwargs):

if name in {"any", "all"}:
Expand Down Expand Up @@ -678,60 +712,3 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):
else:
result[mask] = np.nan
return result

@classmethod
def _create_arithmetic_method(cls, op):
op_name = op.__name__

@ops.unpack_zerodim_and_defer(op_name)
def boolean_arithmetic_method(self, other):
mask = None

if isinstance(other, BooleanArray):
other, mask = 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")

# nans propagate
if mask is None:
mask = self._mask
if other is libmissing.NA:
mask |= True
else:
mask = self._mask | mask

if other is libmissing.NA:
# if other is NA, the result will be all NA and we can't run the
# actual op, so we need to choose the resulting dtype manually
if op_name in {"floordiv", "rfloordiv", "mod", "rmod", "pow", "rpow"}:
dtype = "int8"
else:
dtype = "bool"
result = np.zeros(len(self._data), dtype=dtype)
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)

name = f"__{op_name}__"
return set_function_name(boolean_arithmetic_method, name, cls)


BooleanArray._add_logical_ops()
BooleanArray._add_arithmetic_ops()
128 changes: 56 additions & 72 deletions pandas/core/arrays/floating.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@

from pandas._libs import lib, missing as libmissing
from pandas._typing import ArrayLike, DtypeObj
from pandas.compat import set_function_name
from pandas.compat.numpy import function as nv
from pandas.util._decorators import cache_readonly

Expand All @@ -26,9 +25,7 @@
from pandas.core.dtypes.missing import isna

from pandas.core import ops
from pandas.core.arraylike import OpsMixin
from pandas.core.ops import invalid_comparison
from pandas.core.ops.common import unpack_zerodim_and_defer
from pandas.core.tools.numeric import to_numeric

from .masked import BaseMaskedArray, BaseMaskedDtype
Expand Down Expand Up @@ -202,7 +199,7 @@ def coerce_to_array(
return values, mask


class FloatingArray(OpsMixin, BaseMaskedArray):
class FloatingArray(BaseMaskedArray):
"""
Array of floating (optional missing) values.
Expand Down Expand Up @@ -479,83 +476,70 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):

return type(self)(result, mask, copy=False)

@classmethod
def _create_arithmetic_method(cls, op):
op_name = op.__name__

@unpack_zerodim_and_defer(op.__name__)
def floating_arithmetic_method(self, other):
from pandas.arrays import IntegerArray

omask = None
def _arith_method(self, other, op):
from pandas.arrays import IntegerArray

if getattr(other, "ndim", 0) > 1:
raise NotImplementedError("can only perform ops with 1-d structures")
omask = None

if isinstance(other, (IntegerArray, FloatingArray)):
other, omask = other._data, other._mask
if getattr(other, "ndim", 0) > 1:
raise NotImplementedError("can only perform ops with 1-d structures")

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")
if isinstance(other, (IntegerArray, FloatingArray)):
other, omask = other._data, other._mask

else:
if not (is_float(other) or is_integer(other) or other is libmissing.NA):
raise TypeError("can only perform ops with numeric values")
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")

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)
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:
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)

name = f"__{op.__name__}__"
return set_function_name(floating_arithmetic_method, name, cls)
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"),
)

FloatingArray._add_arithmetic_ops()
return self._maybe_mask_result(result, mask, other, op.__name__)


_dtype_docstring = """
Expand Down
Loading

0 comments on commit 9958b5c

Please sign in to comment.