Skip to content

Commit

Permalink
TYP/REF: use _cmp_method in EAs (pandas-dev#36954)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jbrockmendel
jbrockmendel authored Oct 7, 2020
1 parent 3a83b2c commit 397e36c
Show file tree
Hide file tree
Showing 5 changed files with 166 additions and 191 deletions.
74 changes: 33 additions & 41 deletions pandas/core/arrays/boolean.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
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 @@ -202,7 +203,7 @@ def coerce_to_array(
return values, mask


class BooleanArray(BaseMaskedArray):
class BooleanArray(OpsMixin, BaseMaskedArray):
"""
Array of boolean (True/False) data with missing values.
Expand Down Expand Up @@ -603,52 +604,44 @@ def logical_method(self, other):
name = f"__{op.__name__}__"
return set_function_name(logical_method, name, cls)

@classmethod
def _create_comparison_method(cls, op):
@ops.unpack_zerodim_and_defer(op.__name__)
def cmp_method(self, other):
from pandas.arrays import FloatingArray, IntegerArray
def _cmp_method(self, other, op):
from pandas.arrays import FloatingArray, IntegerArray

if isinstance(other, (IntegerArray, FloatingArray)):
return NotImplemented
if isinstance(other, (IntegerArray, FloatingArray)):
return NotImplemented

mask = None
mask = None

if isinstance(other, BooleanArray):
other, mask = other._data, other._mask
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 to compare")
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 to compare")

if other is libmissing.NA:
# numpy does not handle pd.NA well as "other" scalar (it returns
# a scalar False instead of an array)
result = np.zeros_like(self._data)
mask = np.ones_like(self._data)
else:
# numpy will show a DeprecationWarning on invalid elementwise
# comparisons, this will raise in the future
with warnings.catch_warnings():
warnings.filterwarnings("ignore", "elementwise", FutureWarning)
with np.errstate(all="ignore"):
result = op(self._data, other)

# nans propagate
if mask is None:
mask = self._mask.copy()
else:
mask = self._mask | mask
if other is libmissing.NA:
# numpy does not handle pd.NA well as "other" scalar (it returns
# a scalar False instead of an array)
result = np.zeros_like(self._data)
mask = np.ones_like(self._data)
else:
# numpy will show a DeprecationWarning on invalid elementwise
# comparisons, this will raise in the future
with warnings.catch_warnings():
warnings.filterwarnings("ignore", "elementwise", FutureWarning)
with np.errstate(all="ignore"):
result = op(self._data, other)

return BooleanArray(result, mask, copy=False)
# nans propagate
if mask is None:
mask = self._mask.copy()
else:
mask = self._mask | mask

name = f"__{op.__name__}"
return set_function_name(cmp_method, name, cls)
return BooleanArray(result, mask, copy=False)

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

Expand Down Expand Up @@ -741,5 +734,4 @@ def boolean_arithmetic_method(self, other):


BooleanArray._add_logical_ops()
BooleanArray._add_comparison_ops()
BooleanArray._add_arithmetic_ops()
84 changes: 37 additions & 47 deletions pandas/core/arrays/floating.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,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
Expand Down Expand Up @@ -201,7 +202,7 @@ def coerce_to_array(
return values, mask


class FloatingArray(BaseMaskedArray):
class FloatingArray(OpsMixin, BaseMaskedArray):
"""
Array of floating (optional missing) values.
Expand Down Expand Up @@ -398,58 +399,48 @@ def astype(self, dtype, copy: bool = True) -> ArrayLike:
def _values_for_argsort(self) -> np.ndarray:
return self._data

@classmethod
def _create_comparison_method(cls, op):
op_name = op.__name__
def _cmp_method(self, other, op):
from pandas.arrays import BooleanArray, IntegerArray

@unpack_zerodim_and_defer(op.__name__)
def cmp_method(self, other):
from pandas.arrays import BooleanArray, IntegerArray
mask = None

mask = None
if isinstance(other, (BooleanArray, IntegerArray, FloatingArray)):
other, mask = other._data, other._mask

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

elif is_list_like(other):
other = np.asarray(other)
if other.ndim > 1:
raise NotImplementedError(
"can only perform ops with 1-d structures"
)
if other is libmissing.NA:
# numpy does not handle pd.NA well as "other" scalar (it returns
# a scalar False instead of an array)
# This may be fixed by NA.__array_ufunc__. Revisit this check
# once that's implemented.
result = np.zeros(self._data.shape, dtype="bool")
mask = np.ones(self._data.shape, dtype="bool")
else:
with warnings.catch_warnings():
# numpy may show a FutureWarning:
# elementwise comparison failed; returning scalar instead,
# but in the future will perform elementwise comparison
# before returning NotImplemented. We fall back to the correct
# behavior today, so that should be fine to ignore.
warnings.filterwarnings("ignore", "elementwise", FutureWarning)
with np.errstate(all="ignore"):
method = getattr(self._data, f"__{op.__name__}__")
result = method(other)

if other is libmissing.NA:
# numpy does not handle pd.NA well as "other" scalar (it returns
# a scalar False instead of an array)
# This may be fixed by NA.__array_ufunc__. Revisit this check
# once that's implemented.
result = np.zeros(self._data.shape, dtype="bool")
mask = np.ones(self._data.shape, dtype="bool")
else:
with warnings.catch_warnings():
# numpy may show a FutureWarning:
# elementwise comparison failed; returning scalar instead,
# but in the future will perform elementwise comparison
# before returning NotImplemented. We fall back to the correct
# behavior today, so that should be fine to ignore.
warnings.filterwarnings("ignore", "elementwise", FutureWarning)
with np.errstate(all="ignore"):
method = getattr(self._data, f"__{op_name}__")
result = method(other)

if result is NotImplemented:
result = invalid_comparison(self._data, other, op)

# nans propagate
if mask is None:
mask = self._mask.copy()
else:
mask = self._mask | mask
if result is NotImplemented:
result = invalid_comparison(self._data, other, op)

return BooleanArray(result, mask)
# nans propagate
if mask is None:
mask = self._mask.copy()
else:
mask = self._mask | mask

name = f"__{op.__name__}__"
return set_function_name(cmp_method, name, cls)
return BooleanArray(result, mask)

def sum(self, skipna=True, min_count=0, **kwargs):
nv.validate_sum((), kwargs)
Expand Down Expand Up @@ -565,7 +556,6 @@ def floating_arithmetic_method(self, other):


FloatingArray._add_arithmetic_ops()
FloatingArray._add_comparison_ops()


_dtype_docstring = """
Expand Down
90 changes: 40 additions & 50 deletions pandas/core/arrays/integer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,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
Expand Down Expand Up @@ -265,7 +266,7 @@ def coerce_to_array(
return values, mask


class IntegerArray(BaseMaskedArray):
class IntegerArray(OpsMixin, BaseMaskedArray):
"""
Array of integer (optional missing) values.
Expand Down Expand Up @@ -493,60 +494,50 @@ def _values_for_argsort(self) -> np.ndarray:
data[self._mask] = data.min() - 1
return data

@classmethod
def _create_comparison_method(cls, op):
op_name = op.__name__
def _cmp_method(self, other, op):
from pandas.core.arrays import BaseMaskedArray, BooleanArray

@unpack_zerodim_and_defer(op.__name__)
def cmp_method(self, other):
from pandas.core.arrays import BaseMaskedArray, BooleanArray
mask = None

mask = None
if isinstance(other, BaseMaskedArray):
other, mask = other._data, other._mask

if isinstance(other, BaseMaskedArray):
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 to compare")
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 to compare")

if other is libmissing.NA:
# numpy does not handle pd.NA well as "other" scalar (it returns
# a scalar False instead of an array)
# This may be fixed by NA.__array_ufunc__. Revisit this check
# once that's implemented.
result = np.zeros(self._data.shape, dtype="bool")
mask = np.ones(self._data.shape, dtype="bool")
else:
with warnings.catch_warnings():
# numpy may show a FutureWarning:
# elementwise comparison failed; returning scalar instead,
# but in the future will perform elementwise comparison
# before returning NotImplemented. We fall back to the correct
# behavior today, so that should be fine to ignore.
warnings.filterwarnings("ignore", "elementwise", FutureWarning)
with np.errstate(all="ignore"):
method = getattr(self._data, f"__{op.__name__}__")
result = method(other)

if other is libmissing.NA:
# numpy does not handle pd.NA well as "other" scalar (it returns
# a scalar False instead of an array)
# This may be fixed by NA.__array_ufunc__. Revisit this check
# once that's implemented.
result = np.zeros(self._data.shape, dtype="bool")
mask = np.ones(self._data.shape, dtype="bool")
else:
with warnings.catch_warnings():
# numpy may show a FutureWarning:
# elementwise comparison failed; returning scalar instead,
# but in the future will perform elementwise comparison
# before returning NotImplemented. We fall back to the correct
# behavior today, so that should be fine to ignore.
warnings.filterwarnings("ignore", "elementwise", FutureWarning)
with np.errstate(all="ignore"):
method = getattr(self._data, f"__{op_name}__")
result = method(other)

if result is NotImplemented:
result = invalid_comparison(self._data, other, op)

# nans propagate
if mask is None:
mask = self._mask.copy()
else:
mask = self._mask | mask
if result is NotImplemented:
result = invalid_comparison(self._data, other, op)

return BooleanArray(result, mask)
# nans propagate
if mask is None:
mask = self._mask.copy()
else:
mask = self._mask | mask

name = f"__{op.__name__}__"
return set_function_name(cmp_method, name, cls)
return BooleanArray(result, mask)

def sum(self, skipna=True, min_count=0, **kwargs):
nv.validate_sum((), kwargs)
Expand Down Expand Up @@ -669,7 +660,6 @@ def integer_arithmetic_method(self, other):


IntegerArray._add_arithmetic_ops()
IntegerArray._add_comparison_ops()


_dtype_docstring = """
Expand Down
Loading

0 comments on commit 397e36c

Please sign in to comment.