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