extract overflow-checking wrapper fcn, add tests

Author: patrickmckenna

pandas/_libs/ops.pxd

@@ -0,0 +1,4 @@
+from numpy cimport int64_t
+
+
+cpdef int64_t calculate(object op, int64_t a, int64_t b) except? -1

pandas/_libs/ops.pyi

@@ -48,3 +48,4 @@ def maybe_convert_bool(
     *,
     convert_to_masked_nullable: Literal[True],
 ) -> tuple[np.ndarray, np.ndarray]: ...
+def calculate(op, left: int, right: int) -> int: ...

pandas/_libs/ops.pyx

@@ -16,6 +16,7 @@ import numpy as np
 
 from numpy cimport (
     import_array,
+    int64_t,
     ndarray,
     uint8_t,
 )
@@ -308,3 +309,12 @@ def maybe_convert_bool(ndarray[object] arr,
             return (arr, None)
     else:
         return (result.view(np.bool_), None)
+
+
+@cython.overflowcheck(True)
+cpdef int64_t calculate(object op, int64_t a, int64_t b) except? -1:
+    """
+    Calculate op(a, b) and return the result. Raises OverflowError if converting either
+    operand or the result to an int64_t would overflow.
+    """
+    return op(a, b)

pandas/_libs/tslibs/timedeltas.pyx

@@ -29,6 +29,7 @@ from cpython.datetime cimport (
 import_datetime()
 
 
+from pandas._libs cimport ops
 cimport pandas._libs.tslibs.util as util
 from pandas._libs.tslibs.base cimport ABCTimestamp
 from pandas._libs.tslibs.conversion cimport (
@@ -684,23 +685,13 @@ def _op_unary_method(func, name):
     return f
 
 
-@cython.overflowcheck(True)
-cpdef int64_t _calculate(object op, int64_t a, int64_t b) except? -1:
-    """
-    Calculate op(a, b) and return the result. Raises OverflowError if either operand
-    or the result would overflow on conversion to int64_t.
-    """
-    return op(a, b)
-
-
 cpdef int64_t calculate(object op, object a, object b) except? -1:
     """
-    As above, but raises an OutOfBoundsTimedelta.
+    Calculate op(a, b), raising if either operand or the resulting value cannot be
+    safely cast to an int64_t.
     """
-    cdef int64_t int_a, int_b
-
     try:
-        return _calculate(op, a, b)
+        return ops.calculate(op, a, b)
     except OverflowError as ex:
         msg = f"outside allowed range [{TIMEDELTA_MIN_NS}ns, {TIMEDELTA_MAX_NS}ns]"
         raise OutOfBoundsTimedelta(msg) from ex

pandas/tests/libs/test_ops.py

@@ -0,0 +1,53 @@
+import operator
+
+import numpy as np
+import pytest
+
+from pandas._libs import ops
+
+
+@pytest.fixture(name="int_max")
+def fixture_int_max() -> int:
+    return np.iinfo(np.int64).max
+
+
+@pytest.fixture(name="int_min")
+def fixture_int_min() -> int:
+    return np.iinfo(np.int64).min
+
+
+@pytest.fixture(name="overflow_msg")
+def fixture_overflow_msg() -> str:
+    return "Python int too large to convert to C long"
+
+
+def test_raises_for_too_large_arg(int_max: int, overflow_msg: str):
+    with pytest.raises(OverflowError, match=overflow_msg):
+        ops.calculate(operator.add, int_max + 1, 1)
+
+    with pytest.raises(OverflowError, match=overflow_msg):
+        ops.calculate(operator.add, 1, int_max + 1)
+
+
+def test_raises_for_too_small_arg(int_min: int, overflow_msg: str):
+    with pytest.raises(OverflowError, match=overflow_msg):
+        ops.calculate(operator.add, int_min - 1, 1)
+
+    with pytest.raises(OverflowError, match=overflow_msg):
+        ops.calculate(operator.add, 1, int_min - 1)
+
+
+def test_raises_for_too_large_result(int_max: int, overflow_msg: str):
+    with pytest.raises(OverflowError, match=overflow_msg):
+        ops.calculate(operator.add, int_max, 1)
+
+    with pytest.raises(OverflowError, match=overflow_msg):
+        ops.calculate(operator.add, 1, int_max)
+
+
+def test_raises_for_too_small_result(int_min: int, overflow_msg: str):
+    with pytest.raises(OverflowError, match=overflow_msg):
+        ops.calculate(operator.sub, int_min, 1)
+
+    with pytest.raises(OverflowError, match=overflow_msg):
+        ops.calculate(operator.sub, 1, int_min)