Add more overflow tests for timedelta64 operations

pandas/_testing/__init__.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import collections
+from collections import abc
 from datetime import datetime
 from decimal import Decimal
 from functools import wraps
+from inspect import isclass
 import operator
 import os
 import re
@@ -279,6 +281,28 @@ def box_expected(expected, box_cls, transpose=True):
     return expected
 
 
+def wrap_value(value, cls, transpose=False):
+    """
+    If cls is a scalar type, return value as an instance of it, otherwise return value
+    wrapped in the box type indicated by cls.
+
+    Designed to play nicely with box_expected (and the box_with_array fixture).
+    """
+    if isclass(cls) and not issubclass(cls, pd.core.arraylike.OpsMixin):
+        return cls(value)
+
+    if cls in (np.array, np.ndarray):
+        pass
+    elif cls is pd.array or issubclass(cls, ExtensionArray):
+        cls = pd.array
+    elif issubclass(cls, Index):
+        cls = Index
+
+    if not isinstance(value, (abc.Sequence, pd.core.arraylike.OpsMixin, np.ndarray)):
+        value = [value]
+    return box_expected(value, cls, transpose)
+
+
 def to_array(obj):
     """
     Similar to pd.array, but does not cast numpy dtypes to nullable dtypes.

pandas/tests/arithmetic/test_timedelta64.py

@@ -1,9 +1,18 @@
 # Arithmetic tests for DataFrame/Series/Index/Array classes that should
 # behave identically.
+from contextlib import (
+    AbstractContextManager,
+    nullcontext,
+)
 from datetime import (
     datetime,
     timedelta,
 )
+from functools import partial
+from typing import (
+    Type,
+    Union,
+)
 
 import numpy as np
 import pytest
@@ -31,12 +40,33 @@
     Int64Index,
     UInt64Index,
 )
+from pandas.core.arrays import (
+    DatetimeArray,
+    TimedeltaArray,
+)
 from pandas.tests.arithmetic.common import (
     assert_invalid_addsub_type,
     assert_invalid_comparison,
     get_upcast_box,
 )
 
+TD64_BOX_TYPE = Union[TimedeltaArray, TimedeltaIndex, Series, DataFrame]
+TD64_TYPE = Union[Timedelta, TD64_BOX_TYPE]
+DT64_TYPE = Union[Timestamp, DatetimeArray, DatetimeIndex, Series, DataFrame]
+
+
+TD64_OVERFLOW_MSG = "|".join(
+    [
+        "int too big to convert",
+        "Python int too large to convert to C long",
+        "Overflow in int64 addition",
+    ]
+)
+
+
+does_not_raise = nullcontext
+td64_overflow_error = partial(pytest.raises, OverflowError, match=TD64_OVERFLOW_MSG)
+
 
 def assert_dtype(obj, expected_dtype):
     """
@@ -59,6 +89,59 @@ def get_expected_name(box, names):
     return exname
 
 
+def get_result_type(
+    td64_type: Type[TD64_TYPE],
+    dt64_type: Type[DT64_TYPE],
+) -> Type[DT64_TYPE]:
+    """
+    Expected result when adding, subtracting timedelta64-valued box and
+    datetime64-valued box or scalar.
+    """
+    dt64 = tm.wrap_value(Timestamp.now(), dt64_type)
+    td64 = tm.wrap_value(Timedelta(0), td64_type)
+    return type(dt64 + td64)
+
+
+@pytest.fixture(
+    name="td64_type",
+    params=[Timedelta, TimedeltaArray, TimedeltaIndex, Series, DataFrame],
+    scope="module",
+)
+def fixture_td64_type(request) -> Type[TD64_TYPE]:
+    return request.param
+
+
+@pytest.fixture(
+    name="dt64_type",
+    params=[Timestamp, DatetimeArray, DatetimeIndex, Series, DataFrame],
+    scope="module",
+)
+def fixture_dt64_type(request) -> Type[DT64_TYPE]:
+    return request.param
+
+
+@pytest.fixture(name="max_td64")
+def fixture_max_td64(box_with_array) -> TD64_BOX_TYPE:
+    """
+    A 1-elem ExtensionArray/Index/Series, or 2x1 DataFrame, w/ all elements set to
+    Timestamp.max.
+    """
+    return tm.wrap_value(Timedelta.max, box_with_array)
+
+
+@pytest.fixture(
+    name="positive_td64",
+    params=[Timedelta(1), Timedelta(1024), Timedelta.max],
+    ids=["1ns", "1024ns", "td_max"],
+)
+def fixture_positive_td64(request, td64_type: Type[TD64_TYPE]) -> TD64_TYPE:
+    """
+    A scalar, 1-elem ExtensionArray/Index/Series, or 2x1 DataFrame.
+    """
+    value = request.param
+    return tm.wrap_value(value, td64_type)
+
+
 # ------------------------------------------------------------------
 # Timedelta64[ns] dtype Comparisons
 
@@ -687,30 +770,6 @@ def test_tdarr_add_timestamp_nat_masking(self, box_with_array, str_ts):
                 assert res[1] is NaT
 
     def test_tdi_add_overflow(self):
-        # See GH#14068
-        # preliminary test scalar analogue of vectorized tests below
-        # TODO: Make raised error message more informative and test
-        with pytest.raises(OutOfBoundsDatetime, match="10155196800000000000"):
-            pd.to_timedelta(106580, "D") + Timestamp("2000")
-        with pytest.raises(OutOfBoundsDatetime, match="10155196800000000000"):
-            Timestamp("2000") + pd.to_timedelta(106580, "D")
-
-        _NaT = NaT.value + 1
-        msg = "Overflow in int64 addition"
-        with pytest.raises(OverflowError, match=msg):
-            pd.to_timedelta([106580], "D") + Timestamp("2000")
-        with pytest.raises(OverflowError, match=msg):
-            Timestamp("2000") + pd.to_timedelta([106580], "D")
-        with pytest.raises(OverflowError, match=msg):
-            pd.to_timedelta([_NaT]) - Timedelta("1 days")
-        with pytest.raises(OverflowError, match=msg):
-            pd.to_timedelta(["5 days", _NaT]) - Timedelta("1 days")
-        with pytest.raises(OverflowError, match=msg):
-            (
-                pd.to_timedelta([_NaT, "5 days", "1 hours"])
-                - pd.to_timedelta(["7 seconds", _NaT, "4 hours"])
-            )
-
         # These should not overflow!
         exp = TimedeltaIndex([NaT])
         result = pd.to_timedelta([NaT]) - Timedelta("1 days")
@@ -2089,18 +2148,104 @@ def test_td64arr_pow_invalid(self, scalar_td, box_with_array):
             td1**scalar_td
 
 
-def test_add_timestamp_to_timedelta():
-    # GH: 35897
-    timestamp = Timestamp("2021-01-01")
-    result = timestamp + timedelta_range("0s", "1s", periods=31)
-    expected = DatetimeIndex(
+class TestAddSub:
+    """
+    Addition/subtraction between a timedelta64-valued
+    ExtensionArrays/Indexes/Series/DataFrames, and a timedelta64 scalar or
+    timedelta64-valued ExtensionArray/Index/Series/DataFrame.
+    """
+
+    def test_add_raises_if_result_would_overflow(
+        self,
+        max_td64: TD64_TYPE,
+        positive_td64: TD64_BOX_TYPE,
+    ):
+        with td64_overflow_error():
+            max_td64 + positive_td64
+
+        with td64_overflow_error():
+            positive_td64 + max_td64
+
+    @pytest.mark.parametrize(
+        ["rval", "expected_exs"],
         [
-            timestamp
-            + (
-                pd.to_timedelta("0.033333333s") * i
-                + pd.to_timedelta("0.000000001s") * divmod(i, 3)[0]
-            )
-            for i in range(31)
-        ]
+            (Timedelta(1), does_not_raise()),
+            (Timedelta(2), td64_overflow_error()),
+            (Timedelta.max, td64_overflow_error()),
+        ],
     )
-    tm.assert_index_equal(result, expected)
+    def test_sub_raises_if_result_would_overflow(
+        self,
+        max_td64: TD64_BOX_TYPE,
+        rval: Timedelta,
+        expected_exs: AbstractContextManager,
+        td64_type: Type[TD64_TYPE],
+    ):
+        rvalue = tm.wrap_value(rval, td64_type)
+        min_td64 = -1 * max_td64
+
+        with expected_exs:
+            min_td64 - rvalue
+
+        with expected_exs:
+            -1 * rvalue - max_td64
+
+
+class TestNumericScalarMulDiv:
+    """
+    Operations on timedelta64-valued ExtensionArray/Index/Series/DataFrame and a
+    numeric scalar.
+    """
+
+    @pytest.mark.xfail(reason="Not implemented")
+    def test_scalar_mul_raises_if_result_would_overflow(self, max_td64: TD64_BOX_TYPE):
+        with td64_overflow_error():
+            max_td64 * 1.01
+
+        with td64_overflow_error():
+            1.01 * max_td64
+
+
+class TestAddSubDatetime64:
+    """
+    Operations on timedelta64-valued ExtensionArray/Index/Series/DataFrame, and a
+    datetime64 scalar or datetime64-valued ExtensionArray/Index/Series/DataFrame.
+    """
+
+    def test_add(self, box_with_array, dt64_type: Type[DT64_TYPE]):
+        # GH: 35897
+        dt64 = tm.wrap_value(Timestamp(2020, 1, 2), dt64_type)
+        td64_box = tm.wrap_value(Timedelta(hours=3), box_with_array)
+
+        expected_type = get_result_type(box_with_array, dt64_type)
+        expected = tm.wrap_value(Timestamp(2020, 1, 2, 3), expected_type)
+        result = dt64 + td64_box
+
+        tm.assert_equal(result, expected)
+
+    def test_add_dt64_raises_if_result_would_overflow(
+        self,
+        max_td64: TD64_BOX_TYPE,
+        dt64_type: Type[DT64_TYPE],
+    ):
+        max_dt64 = tm.wrap_value(Timestamp.max, dt64_type)
+        ex = (OutOfBoundsDatetime, OverflowError)
+        msg = "|".join([TD64_OVERFLOW_MSG, "Out of bounds nanosecond timestamp"])
+
+        with pytest.raises(ex, match=msg):
+            max_td64 + max_dt64
+
+        with pytest.raises(ex, match=msg):
+            max_dt64 + max_td64
+
+    def test_sub_td64_raises_if_result_would_overflow(
+        self,
+        max_td64: TD64_BOX_TYPE,
+        dt64_type: Type[DT64_TYPE],
+    ):
+        min_dt64 = tm.wrap_value(Timestamp.min, dt64_type)
+        ex = (OutOfBoundsDatetime, OverflowError)
+        msg = "|".join([TD64_OVERFLOW_MSG, "Out of bounds nanosecond timestamp"])
+
+        with pytest.raises(ex, match=msg):
+            min_dt64 - max_td64