Skip to content

REF: use OpsMixin in EAs #37049

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Oct 11, 2020
Merged

REF: use OpsMixin in EAs #37049

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d68eab5
REF: use OpsMixin in EAs
jbrockmendel Oct 11, 2020
499cd13
Merge branch 'master' of https://github.com/pandas-dev/pandas into re…
jbrockmendel Oct 11, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
175 changes: 76 additions & 99 deletions pandas/core/arrays/boolean.py
View file
Open in desktop
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
Open in desktop
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