TST: xfailed arithmetic tests (pandas-dev#36870)

Author: jbrockmendel

doc/source/whatsnew/v1.2.0.rst

@@ -336,6 +336,7 @@ Numeric
 - Bug in :class:`Series` where two :class:`Series` each have a :class:`DatetimeIndex` with different timezones having those indexes incorrectly changed when performing arithmetic operations (:issue:`33671`)
 - Bug in :meth:`pd._testing.assert_almost_equal` was incorrect for complex numeric types (:issue:`28235`)
 - Bug in :meth:`DataFrame.__rmatmul__` error handling reporting transposed shapes (:issue:`21581`)
+- Bug in :class:`IntegerArray` multiplication with ``timedelta`` and ``np.timedelta64`` objects (:issue:`36870`)
 
 Conversion
 ^^^^^^^^^^

pandas/core/arrays/integer.py

@@ -1,10 +1,11 @@
+from datetime import timedelta
 import numbers
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
 import warnings
 
 import numpy as np
 
-from pandas._libs import lib, missing as libmissing
+from pandas._libs import Timedelta, iNaT, 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
@@ -581,6 +582,12 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):
             result[mask] = np.nan
             return result
 
+        if result.dtype == "timedelta64[ns]":
+            from pandas.core.arrays import TimedeltaArray
+
+            result[mask] = iNaT
+            return TimedeltaArray._simple_new(result)
+
         return type(self)(result, mask, copy=False)
 
     @classmethod
@@ -609,6 +616,9 @@ def integer_arithmetic_method(self, other):
                 if not (is_float_dtype(other) or is_integer_dtype(other)):
                     raise TypeError("can only perform ops with numeric values")
 
+            elif isinstance(other, (timedelta, np.timedelta64)):
+                other = Timedelta(other)
+
             else:
                 if not (is_float(other) or is_integer(other) or other is libmissing.NA):
                     raise TypeError("can only perform ops with numeric values")

pandas/core/arrays/numpy_.py

@@ -218,6 +218,16 @@ def __array_ufunc__(self, ufunc, method: str, *inputs, **kwargs):
             if not isinstance(x, self._HANDLED_TYPES + (PandasArray,)):
                 return NotImplemented
 
+        if ufunc not in [np.logical_or, np.bitwise_or, np.bitwise_xor]:
+            # For binary ops, use our custom dunder methods
+            # We haven't implemented logical dunder funcs, so exclude these
+            #  to avoid RecursionError
+            result = ops.maybe_dispatch_ufunc_to_dunder_op(
+                self, ufunc, method, *inputs, **kwargs
+            )
+            if result is not NotImplemented:
+                return result
+
         # Defer to the implementation of the ufunc on unwrapped values.
         inputs = tuple(x._ndarray if isinstance(x, PandasArray) else x for x in inputs)
         if out:
@@ -377,19 +387,28 @@ def arithmetic_method(self, other):
                 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 self._wrap_ndarray_result(a), self._wrap_ndarray_result(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 self._wrap_ndarray_result(result)
             return result
 
         return compat.set_function_name(arithmetic_method, f"__{op.__name__}__", cls)
 
     _create_comparison_method = _create_arithmetic_method
 
+    def _wrap_ndarray_result(self, result: np.ndarray):
+        # If we have timedelta64[ns] result, return a TimedeltaArray instead
+        #  of a PandasArray
+        if result.dtype == "timedelta64[ns]":
+            from pandas.core.arrays import TimedeltaArray
+
+            return TimedeltaArray._simple_new(result)
+        return type(self)(result)
+
     # ------------------------------------------------------------------------
     # String methods interface
     _str_na_value = np.nan

pandas/tests/arithmetic/common.py

@@ -6,6 +6,7 @@
 
 from pandas import DataFrame, Index, Series, array as pd_array
 import pandas._testing as tm
+from pandas.core.arrays import PandasArray
 
 
 def assert_invalid_addsub_type(left, right, msg=None):
@@ -56,18 +57,25 @@ def assert_invalid_comparison(left, right, box):
     # Note: not quite the same as how we do this for tm.box_expected
     xbox = box if box not in [Index, pd_array] else np.array
 
-    result = left == right
+    def xbox2(x):
+        # Eventually we'd like this to be tighter, but for now we'll
+        #  just exclude PandasArray[bool]
+        if isinstance(x, PandasArray):
+            return x._ndarray
+        return x
+
+    result = xbox2(left == right)
     expected = xbox(np.zeros(result.shape, dtype=np.bool_))
 
     tm.assert_equal(result, expected)
 
-    result = right == left
+    result = xbox2(right == left)
     tm.assert_equal(result, expected)
 
-    result = left != right
+    result = xbox2(left != right)
     tm.assert_equal(result, ~expected)
 
-    result = right != left
+    result = xbox2(right != left)
     tm.assert_equal(result, ~expected)
 
     msg = "|".join(

pandas/tests/arithmetic/conftest.py

@@ -221,15 +221,6 @@ def mismatched_freq(request):
 # ------------------------------------------------------------------
 
 
-@pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame], ids=id_func)
-def box(request):
-    """
-    Several array-like containers that should have effectively identical
-    behavior with respect to arithmetic operations.
-    """
-    return request.param
-
-
 @pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, pd.array], ids=id_func)
 def box_with_array(request):
     """

pandas/tests/arithmetic/test_datetime64.py

@@ -205,8 +205,6 @@ 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,32 +224,36 @@ def test_comparison_invalid(self, tz_naive_fixture, box_with_array):
     )
     @pytest.mark.parametrize("dtype", [None, object])
     def test_nat_comparisons_scalar(self, dtype, data, box_with_array):
+        box = box_with_array
         if box_with_array is tm.to_array and dtype is object:
             # 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 not in [pd.Index, pd.array] else np.ndarray
-        )
+        xbox = box if box not in [pd.Index, pd.array] else np.ndarray
 
         left = Series(data, dtype=dtype)
-        left = tm.box_expected(left, box_with_array)
+        left = tm.box_expected(left, box)
 
         expected = [False, False, False]
         expected = tm.box_expected(expected, xbox)
+        if box is pd.array and dtype is object:
+            expected = pd.array(expected, dtype="bool")
+
         tm.assert_equal(left == NaT, expected)
         tm.assert_equal(NaT == left, expected)
 
         expected = [True, True, True]
         expected = tm.box_expected(expected, xbox)
+        if box is pd.array and dtype is object:
+            expected = pd.array(expected, dtype="bool")
         tm.assert_equal(left != NaT, expected)
         tm.assert_equal(NaT != left, expected)
 
         expected = [False, False, False]
         expected = tm.box_expected(expected, xbox)
+        if box is pd.array and dtype is object:
+            expected = pd.array(expected, dtype="bool")
         tm.assert_equal(left < NaT, expected)
         tm.assert_equal(NaT > left, expected)
         tm.assert_equal(left <= NaT, expected)

pandas/tests/arithmetic/test_numeric.py

@@ -187,10 +187,6 @@ def test_ops_series(self):
     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)])
 
@@ -214,8 +210,6 @@ def test_numeric_arr_mul_tdscalar(self, scalar_td, numeric_idx, box_with_array):
     )
     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)
@@ -231,8 +225,6 @@ def test_numeric_arr_mul_tdscalar_numexpr_path(self, scalar_td, box_with_array):
 
     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]
 
@@ -263,8 +255,6 @@ def test_numeric_arr_rdiv_tdscalar(self, three_days, numeric_idx, box_with_array
     )
     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 = (