ENH: PandasArray ops use core.ops functions (#36484)

Author: jbrockmendel

pandas/core/arrays/numpy_.py

@@ -362,19 +362,29 @@ def __invert__(self):
 
     @classmethod
     def _create_arithmetic_method(cls, op):
+
+        pd_op = ops.get_array_op(op)
+
         @ops.unpack_zerodim_and_defer(op.__name__)
         def arithmetic_method(self, other):
             if isinstance(other, cls):
                 other = other._ndarray
 
-            with np.errstate(all="ignore"):
-                result = op(self._ndarray, other)
+            result = pd_op(self._ndarray, other)
 
-            if op is divmod:
+            if op is divmod or op is ops.rdivmod:
                 a, b = result
-                return cls(a), cls(b)
-
-            return cls(result)
+                if isinstance(a, np.ndarray):
+                    # for e.g. op vs TimedeltaArray, we may already
+                    #  have an ExtensionArray, in which case we do not wrap
+                    return cls(a), cls(b)
+                return a, b
+
+            if isinstance(result, np.ndarray):
+                # for e.g. multiplication vs TimedeltaArray, we may already
+                #  have an ExtensionArray, in which case we do not wrap
+                return cls(result)
+            return result
 
         return compat.set_function_name(arithmetic_method, f"__{op.__name__}__", cls)
 

pandas/tests/arithmetic/common.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pytest
 
-from pandas import DataFrame, Index, Series
+from pandas import DataFrame, Index, Series, array as pd_array
 import pandas._testing as tm
 
 
@@ -49,12 +49,12 @@ def assert_invalid_comparison(left, right, box):
     ----------
     left : np.ndarray, ExtensionArray, Index, or Series
     right : object
-    box : {pd.DataFrame, pd.Series, pd.Index, tm.to_array}
+    box : {pd.DataFrame, pd.Series, pd.Index, pd.array, tm.to_array}
     """
     # Not for tznaive-tzaware comparison
 
     # Note: not quite the same as how we do this for tm.box_expected
-    xbox = box if box is not Index else np.array
+    xbox = box if box not in [Index, pd_array] else np.array
 
     result = left == right
     expected = xbox(np.zeros(result.shape, dtype=np.bool_))

pandas/tests/arithmetic/conftest.py

@@ -2,7 +2,6 @@
 import pytest
 
 import pandas as pd
-import pandas._testing as tm
 
 # ------------------------------------------------------------------
 # Helper Functions
@@ -56,7 +55,7 @@ def one(request):
 
 zeros = [
     box_cls([0] * 5, dtype=dtype)
-    for box_cls in [pd.Index, np.array]
+    for box_cls in [pd.Index, np.array, pd.array]
     for dtype in [np.int64, np.uint64, np.float64]
 ]
 zeros.extend(
@@ -231,7 +230,7 @@ def box(request):
     return request.param
 
 
-@pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, tm.to_array], ids=id_func)
+@pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, pd.array], ids=id_func)
 def box_with_array(request):
     """
     Fixture to test behavior for Index, Series, DataFrame, and pandas Array

pandas/tests/arithmetic/test_datetime64.py

@@ -48,7 +48,9 @@ def test_compare_zerodim(self, tz_naive_fixture, box_with_array):
         # Test comparison with zero-dimensional array is unboxed
         tz = tz_naive_fixture
         box = box_with_array
-        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+        xbox = (
+            box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray
+        )
         dti = date_range("20130101", periods=3, tz=tz)
 
         other = np.array(dti.to_numpy()[0])
@@ -135,7 +137,7 @@ def test_dt64arr_nat_comparison(self, tz_naive_fixture, box_with_array):
         # GH#22242, GH#22163 DataFrame considered NaT == ts incorrectly
         tz = tz_naive_fixture
         box = box_with_array
-        xbox = box if box is not pd.Index else np.ndarray
+        xbox = box if box not in [pd.Index, pd.array] else np.ndarray
 
         ts = pd.Timestamp.now(tz)
         ser = pd.Series([ts, pd.NaT])
@@ -203,6 +205,8 @@ def test_nat_comparisons(self, dtype, index_or_series, reverse, pair):
     def test_comparison_invalid(self, tz_naive_fixture, box_with_array):
         # GH#4968
         # invalid date/int comparisons
+        if box_with_array is pd.array:
+            pytest.xfail("assert_invalid_comparison doesnt handle BooleanArray yet")
         tz = tz_naive_fixture
         ser = Series(range(5))
         ser2 = Series(pd.date_range("20010101", periods=5, tz=tz))
@@ -226,8 +230,12 @@ def test_nat_comparisons_scalar(self, dtype, data, box_with_array):
             # dont bother testing ndarray comparison methods as this fails
             #  on older numpys (since they check object identity)
             return
+        if box_with_array is pd.array and dtype is object:
+            pytest.xfail("reversed comparisons give BooleanArray, not ndarray")
 
-        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+        xbox = (
+            box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray
+        )
 
         left = Series(data, dtype=dtype)
         left = tm.box_expected(left, box_with_array)
@@ -299,7 +307,9 @@ def test_timestamp_compare_series(self, left, right):
 
     def test_dt64arr_timestamp_equality(self, box_with_array):
         # GH#11034
-        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+        xbox = (
+            box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray
+        )
 
         ser = pd.Series([pd.Timestamp("2000-01-29 01:59:00"), "NaT"])
         ser = tm.box_expected(ser, box_with_array)
@@ -388,7 +398,9 @@ def test_dti_cmp_nat(self, dtype, box_with_array):
             #  on older numpys (since they check object identity)
             return
 
-        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+        xbox = (
+            box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray
+        )
 
         left = pd.DatetimeIndex(
             [pd.Timestamp("2011-01-01"), pd.NaT, pd.Timestamp("2011-01-03")]

pandas/tests/arithmetic/test_numeric.py

@@ -89,8 +89,9 @@ def test_compare_invalid(self):
         b.name = pd.Timestamp("2000-01-01")
         tm.assert_series_equal(a / b, 1 / (b / a))
 
-    def test_numeric_cmp_string_numexpr_path(self, box):
+    def test_numeric_cmp_string_numexpr_path(self, box_with_array):
         # GH#36377, GH#35700
+        box = box_with_array
         xbox = box if box is not pd.Index else np.ndarray
 
         obj = pd.Series(np.random.randn(10 ** 5))
@@ -183,10 +184,14 @@ def test_ops_series(self):
         ],
         ids=lambda x: type(x).__name__,
     )
-    def test_numeric_arr_mul_tdscalar(self, scalar_td, numeric_idx, box):
+    def test_numeric_arr_mul_tdscalar(self, scalar_td, numeric_idx, box_with_array):
         # GH#19333
+        box = box_with_array
+        if box is pd.array:
+            pytest.xfail(
+                "we get a PandasArray[timedelta64[ns]] instead of TimedeltaArray"
+            )
         index = numeric_idx
-
         expected = pd.TimedeltaIndex([pd.Timedelta(days=n) for n in range(5)])
 
         index = tm.box_expected(index, box)
@@ -207,7 +212,11 @@ def test_numeric_arr_mul_tdscalar(self, scalar_td, numeric_idx, box):
         ],
         ids=lambda x: type(x).__name__,
     )
-    def test_numeric_arr_mul_tdscalar_numexpr_path(self, scalar_td, box):
+    def test_numeric_arr_mul_tdscalar_numexpr_path(self, scalar_td, box_with_array):
+        box = box_with_array
+        if box is pd.array:
+            pytest.xfail("IntegerArray.__mul__ doesnt handle timedeltas")
+
         arr = np.arange(2 * 10 ** 4).astype(np.int64)
         obj = tm.box_expected(arr, box, transpose=False)
 
@@ -220,7 +229,11 @@ def test_numeric_arr_mul_tdscalar_numexpr_path(self, scalar_td, box):
         result = scalar_td * obj
         tm.assert_equal(result, expected)
 
-    def test_numeric_arr_rdiv_tdscalar(self, three_days, numeric_idx, box):
+    def test_numeric_arr_rdiv_tdscalar(self, three_days, numeric_idx, box_with_array):
+        box = box_with_array
+        if box is pd.array:
+            pytest.xfail("We get PandasArray[td64] instead of TimedeltaArray")
+
         index = numeric_idx[1:3]
 
         expected = TimedeltaIndex(["3 Days", "36 Hours"])
@@ -248,7 +261,11 @@ def test_numeric_arr_rdiv_tdscalar(self, three_days, numeric_idx, box):
             pd.offsets.Second(0),
         ],
     )
-    def test_add_sub_timedeltalike_invalid(self, numeric_idx, other, box):
+    def test_add_sub_timedeltalike_invalid(self, numeric_idx, other, box_with_array):
+        box = box_with_array
+        if box is pd.array:
+            pytest.xfail("PandasArray[int].__add__ doesnt raise on td64")
+
         left = tm.box_expected(numeric_idx, box)
         msg = (
             "unsupported operand type|"
@@ -276,16 +293,21 @@ def test_add_sub_timedeltalike_invalid(self, numeric_idx, other, box):
         ],
     )
     @pytest.mark.filterwarnings("ignore:elementwise comp:DeprecationWarning")
-    def test_add_sub_datetimelike_invalid(self, numeric_idx, other, box):
+    def test_add_sub_datetimelike_invalid(self, numeric_idx, other, box_with_array):
         # GH#28080 numeric+datetime64 should raise; Timestamp raises
         #  NullFrequencyError instead of TypeError so is excluded.
+        box = box_with_array
         left = tm.box_expected(numeric_idx, box)
 
-        msg = (
-            "unsupported operand type|"
-            "Cannot (add|subtract) NaT (to|from) ndarray|"
-            "Addition/subtraction of integers and integer-arrays|"
-            "Concatenation operation is not implemented for NumPy arrays"
+        msg = "|".join(
+            [
+                "unsupported operand type",
+                "Cannot (add|subtract) NaT (to|from) ndarray",
+                "Addition/subtraction of integers and integer-arrays",
+                "Concatenation operation is not implemented for NumPy arrays",
+                # pd.array vs np.datetime64 case
+                r"operand type\(s\) all returned NotImplemented from __array_ufunc__",
+            ]
         )
         with pytest.raises(TypeError, match=msg):
             left + other
@@ -568,8 +590,9 @@ class TestMultiplicationDivision:
     # __mul__, __rmul__, __div__, __rdiv__, __floordiv__, __rfloordiv__
     # for non-timestamp/timedelta/period dtypes
 
-    def test_divide_decimal(self, box):
+    def test_divide_decimal(self, box_with_array):
         # resolves issue GH#9787
+        box = box_with_array
         ser = Series([Decimal(10)])
         expected = Series([Decimal(5)])
 

pandas/tests/arithmetic/test_object.py

@@ -104,22 +104,22 @@ def test_add_extension_scalar(self, other, box_with_array, op):
         result = op(arr, other)
         tm.assert_equal(result, expected)
 
-    def test_objarr_add_str(self, box):
+    def test_objarr_add_str(self, box_with_array):
         ser = pd.Series(["x", np.nan, "x"])
         expected = pd.Series(["xa", np.nan, "xa"])
 
-        ser = tm.box_expected(ser, box)
-        expected = tm.box_expected(expected, box)
+        ser = tm.box_expected(ser, box_with_array)
+        expected = tm.box_expected(expected, box_with_array)
 
         result = ser + "a"
         tm.assert_equal(result, expected)
 
-    def test_objarr_radd_str(self, box):
+    def test_objarr_radd_str(self, box_with_array):
         ser = pd.Series(["x", np.nan, "x"])
         expected = pd.Series(["ax", np.nan, "ax"])
 
-        ser = tm.box_expected(ser, box)
-        expected = tm.box_expected(expected, box)
+        ser = tm.box_expected(ser, box_with_array)
+        expected = tm.box_expected(expected, box_with_array)
 
         result = "a" + ser
         tm.assert_equal(result, expected)

pandas/tests/arithmetic/test_period.py

@@ -28,7 +28,9 @@ class TestPeriodArrayLikeComparisons:
 
     def test_compare_zerodim(self, box_with_array):
         # GH#26689 make sure we unbox zero-dimensional arrays
-        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+        xbox = (
+            box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray
+        )
 
         pi = pd.period_range("2000", periods=4)
         other = np.array(pi.to_numpy()[0])
@@ -68,7 +70,7 @@ def test_compare_object_dtype(self, box_with_array, other_box):
         pi = pd.period_range("2000", periods=5)
         parr = tm.box_expected(pi, box_with_array)
 
-        xbox = np.ndarray if box_with_array is pd.Index else box_with_array
+        xbox = np.ndarray if box_with_array in [pd.Index, pd.array] else box_with_array
 
         other = other_box(pi)
 
@@ -175,7 +177,9 @@ def test_pi_cmp_period(self):
 
     # TODO: moved from test_datetime64; de-duplicate with version below
     def test_parr_cmp_period_scalar2(self, box_with_array):
-        xbox = box_with_array if box_with_array is not pd.Index else np.ndarray
+        xbox = (
+            box_with_array if box_with_array not in [pd.Index, pd.array] else np.ndarray
+        )
 
         pi = pd.period_range("2000-01-01", periods=10, freq="D")
 
@@ -196,7 +200,7 @@ def test_parr_cmp_period_scalar2(self, box_with_array):
     @pytest.mark.parametrize("freq", ["M", "2M", "3M"])
     def test_parr_cmp_period_scalar(self, freq, box_with_array):
         # GH#13200
-        xbox = np.ndarray if box_with_array is pd.Index else box_with_array
+        xbox = np.ndarray if box_with_array in [pd.Index, pd.array] else box_with_array
 
         base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq)
         base = tm.box_expected(base, box_with_array)
@@ -235,7 +239,7 @@ def test_parr_cmp_period_scalar(self, freq, box_with_array):
     @pytest.mark.parametrize("freq", ["M", "2M", "3M"])
     def test_parr_cmp_pi(self, freq, box_with_array):
         # GH#13200
-        xbox = np.ndarray if box_with_array is pd.Index else box_with_array
+        xbox = np.ndarray if box_with_array in [pd.Index, pd.array] else box_with_array
 
         base = PeriodIndex(["2011-01", "2011-02", "2011-03", "2011-04"], freq=freq)
         base = tm.box_expected(base, box_with_array)
@@ -284,7 +288,7 @@ def test_parr_cmp_pi_mismatched_freq_raises(self, freq, box_with_array):
         # TODO: Could parametrize over boxes for idx?
         idx = PeriodIndex(["2011", "2012", "2013", "2014"], freq="A")
         rev_msg = r"Input has different freq=(M|2M|3M) from PeriodArray\(freq=A-DEC\)"
-        idx_msg = rev_msg if box_with_array is tm.to_array else msg
+        idx_msg = rev_msg if box_with_array in [tm.to_array, pd.array] else msg
         with pytest.raises(IncompatibleFrequency, match=idx_msg):
             base <= idx
 
@@ -298,7 +302,7 @@ def test_parr_cmp_pi_mismatched_freq_raises(self, freq, box_with_array):
 
         idx = PeriodIndex(["2011", "2012", "2013", "2014"], freq="4M")
         rev_msg = r"Input has different freq=(M|2M|3M) from PeriodArray\(freq=4M\)"
-        idx_msg = rev_msg if box_with_array is tm.to_array else msg
+        idx_msg = rev_msg if box_with_array in [tm.to_array, pd.array] else msg
         with pytest.raises(IncompatibleFrequency, match=idx_msg):
             base <= idx
 
@@ -779,7 +783,7 @@ def test_pi_add_sub_td64_array_tick(self):
     @pytest.mark.parametrize("tdi_freq", [None, "H"])
     def test_parr_sub_td64array(self, box_with_array, tdi_freq, pi_freq):
         box = box_with_array
-        xbox = box if box is not tm.to_array else pd.Index
+        xbox = box if box not in [pd.array, tm.to_array] else pd.Index
 
         tdi = TimedeltaIndex(["1 hours", "2 hours"], freq=tdi_freq)
         dti = Timestamp("2018-03-07 17:16:40") + tdi