REF: use unpack_zerodim_and_defer on EA methods (#34042)

Author: jbrockmendel

pandas/core/arrays/base.py

@@ -22,7 +22,7 @@
 from pandas.core.dtypes.cast import maybe_cast_to_extension_array
 from pandas.core.dtypes.common import is_array_like, is_list_like, pandas_dtype
 from pandas.core.dtypes.dtypes import ExtensionDtype
-from pandas.core.dtypes.generic import ABCIndexClass, ABCSeries
+from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
@@ -1273,7 +1273,7 @@ def convert_values(param):
                     ovalues = [param] * len(self)
                 return ovalues
 
-            if isinstance(other, (ABCSeries, ABCIndexClass)):
+            if isinstance(other, (ABCSeries, ABCIndexClass, ABCDataFrame)):
                 # rely on pandas to unbox and dispatch to us
                 return NotImplemented
 

pandas/core/arrays/boolean.py

@@ -20,7 +20,6 @@
     pandas_dtype,
 )
 from pandas.core.dtypes.dtypes import register_extension_dtype
-from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
@@ -559,13 +558,10 @@ def all(self, skipna: bool = True, **kwargs):
 
     @classmethod
     def _create_logical_method(cls, op):
+        @ops.unpack_zerodim_and_defer(op.__name__)
         def logical_method(self, other):
-            if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
-                # Rely on pandas to unbox and dispatch to us.
-                return NotImplemented
 
             assert op.__name__ in {"or_", "ror_", "and_", "rand_", "xor", "rxor"}
-            other = lib.item_from_zerodim(other)
             other_is_booleanarray = isinstance(other, BooleanArray)
             other_is_scalar = lib.is_scalar(other)
             mask = None
@@ -605,16 +601,14 @@ def logical_method(self, other):
 
     @classmethod
     def _create_comparison_method(cls, op):
+        @ops.unpack_zerodim_and_defer(op.__name__)
         def cmp_method(self, other):
             from pandas.arrays import IntegerArray
 
-            if isinstance(
-                other, (ABCDataFrame, ABCSeries, ABCIndexClass, IntegerArray)
-            ):
+            if isinstance(other, IntegerArray):
                 # Rely on pandas to unbox and dispatch to us.
                 return NotImplemented
 
-            other = lib.item_from_zerodim(other)
             mask = None
 
             if isinstance(other, BooleanArray):
@@ -693,13 +687,8 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):
     def _create_arithmetic_method(cls, op):
         op_name = op.__name__
 
+        @ops.unpack_zerodim_and_defer(op_name)
         def boolean_arithmetic_method(self, other):
-
-            if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
-                # Rely on pandas to unbox and dispatch to us.
-                return NotImplemented
-
-            other = lib.item_from_zerodim(other)
             mask = None
 
             if isinstance(other, BooleanArray):

pandas/core/arrays/numpy_.py

@@ -11,12 +11,11 @@
 from pandas.util._validators import validate_fillna_kwargs
 
 from pandas.core.dtypes.dtypes import ExtensionDtype
-from pandas.core.dtypes.generic import ABCIndexClass, ABCSeries
 from pandas.core.dtypes.inference import is_array_like
 from pandas.core.dtypes.missing import isna
 
 from pandas import compat
-from pandas.core import nanops
+from pandas.core import nanops, ops
 from pandas.core.algorithms import searchsorted
 from pandas.core.array_algos import masked_reductions
 from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
@@ -436,11 +435,9 @@ def __invert__(self):
 
     @classmethod
     def _create_arithmetic_method(cls, op):
+        @ops.unpack_zerodim_and_defer(op.__name__)
         def arithmetic_method(self, other):
-            if isinstance(other, (ABCIndexClass, ABCSeries)):
-                return NotImplemented
-
-            elif isinstance(other, cls):
+            if isinstance(other, cls):
                 other = other._ndarray
 
             with np.errstate(all="ignore"):

pandas/core/arrays/string_.py

@@ -7,7 +7,6 @@
 
 from pandas.core.dtypes.base import ExtensionDtype, register_extension_dtype
 from pandas.core.dtypes.common import pandas_dtype
-from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
 from pandas.core.dtypes.inference import is_array_like
 
 from pandas import compat
@@ -312,15 +311,14 @@ def memory_usage(self, deep=False):
     @classmethod
     def _create_arithmetic_method(cls, op):
         # Note: this handles both arithmetic and comparison methods.
+
+        @ops.unpack_zerodim_and_defer(op.__name__)
         def method(self, other):
             from pandas.arrays import BooleanArray
 
             assert op.__name__ in ops.ARITHMETIC_BINOPS | ops.COMPARISON_BINOPS
 
-            if isinstance(other, (ABCIndexClass, ABCSeries, ABCDataFrame)):
-                return NotImplemented
-
-            elif isinstance(other, cls):
+            if isinstance(other, cls):
                 other = other._ndarray
 
             mask = isna(self) | isna(other)

pandas/tests/extension/base/ops.py

@@ -114,10 +114,13 @@ def test_error(self, data, all_arithmetic_operators):
         with pytest.raises(AttributeError):
             getattr(data, op_name)
 
-    def test_direct_arith_with_series_returns_not_implemented(self, data):
-        # EAs should return NotImplemented for ops with Series.
+    @pytest.mark.parametrize("box", [pd.Series, pd.DataFrame])
+    def test_direct_arith_with_ndframe_returns_not_implemented(self, data, box):
+        # EAs should return NotImplemented for ops with Series/DataFrame
         # Pandas takes care of unboxing the series and calling the EA's op.
         other = pd.Series(data)
+        if box is pd.DataFrame:
+            other = other.to_frame()
         if hasattr(data, "__add__"):
             result = data.__add__(other)
             assert result is NotImplemented
@@ -156,10 +159,14 @@ def test_compare_array(self, data, all_compare_operators):
         other = pd.Series([data[0]] * len(data))
         self._compare_other(s, data, op_name, other)
 
-    def test_direct_arith_with_series_returns_not_implemented(self, data):
-        # EAs should return NotImplemented for ops with Series.
+    @pytest.mark.parametrize("box", [pd.Series, pd.DataFrame])
+    def test_direct_arith_with_ndframe_returns_not_implemented(self, data, box):
+        # EAs should return NotImplemented for ops with Series/DataFrame
         # Pandas takes care of unboxing the series and calling the EA's op.
         other = pd.Series(data)
+        if box is pd.DataFrame:
+            other = other.to_frame()
+
         if hasattr(data, "__eq__"):
             result = data.__eq__(other)
             assert result is NotImplemented

pandas/tests/extension/test_period.py

@@ -126,9 +126,13 @@ def test_add_series_with_extension_array(self, data):
     def test_error(self):
         pass
 
-    def test_direct_arith_with_series_returns_not_implemented(self, data):
+    @pytest.mark.parametrize("box", [pd.Series, pd.DataFrame])
+    def test_direct_arith_with_ndframe_returns_not_implemented(self, data, box):
         # Override to use __sub__ instead of __add__
         other = pd.Series(data)
+        if box is pd.DataFrame:
+            other = other.to_frame()
+
         result = data.__sub__(other)
         assert result is NotImplemented