add wrapper for overflow-checked int/float ops

Author: patrickmckenna

pandas/_libs/ops.pxd

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

pandas/_libs/ops.pyi

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

pandas/_libs/ops.pyx

@@ -312,9 +312,17 @@ def maybe_convert_bool(ndarray[object] arr,
 
 
 @cython.overflowcheck(True)
-cpdef int64_t calculate(object op, int64_t a, int64_t b) except? -1:
+cpdef int64_t calc_int_int(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)
+
+@cython.overflowcheck(True)
+cpdef int64_t calc_int_float(object op, int64_t a, double b) except? -1:
+    """
+    Calculate op(a, b) and return the result. Raises OverflowError if converting either
+    operand or the result would overflow.
+    """
+    return op(a, b)

pandas/tests/libs/test_ops.py

@@ -6,17 +6,27 @@
 from pandas._libs import ops
 
 
-@pytest.fixture(name="int_max")
+@pytest.fixture(name="int_max", scope="module")
 def fixture_int_max() -> int:
     return np.iinfo(np.int64).max
 
 
-@pytest.fixture(name="int_min")
+@pytest.fixture(name="int_min", scope="module")
 def fixture_int_min() -> int:
     return np.iinfo(np.int64).min
 
 
-@pytest.fixture(name="overflow_msg")
+@pytest.fixture(name="float_max", scope="module")
+def fixture_float_max() -> int:
+    return np.finfo(np.float64).max
+
+
+@pytest.fixture(name="float_min", scope="module")
+def fixture_float_min() -> int:
+    return np.finfo(np.float64).min
+
+
+@pytest.fixture(name="overflow_msg", scope="module")
 def fixture_overflow_msg() -> str:
     return "|".join(
         (
@@ -26,33 +36,108 @@ def fixture_overflow_msg() -> str:
     )
 
 
-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)
+class TestCalcIntInt:
+    def test_raises_for_too_large_arg(self, int_max: int, overflow_msg: str):
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_int(operator.add, int_max + 1, 1)
 
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_int(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)
+    def test_raises_for_too_small_arg(self, int_min: int, overflow_msg: str):
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_int(operator.add, int_min - 1, 1)
 
-    with pytest.raises(OverflowError, match=overflow_msg):
-        ops.calculate(operator.add, 1, int_min - 1)
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_int(operator.add, 1, int_min - 1)
 
+    def test_raises_for_too_large_result(self, int_max: int, overflow_msg: str):
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_int(operator.add, int_max, 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.calc_int_int(operator.add, 1, int_max)
 
-    with pytest.raises(OverflowError, match=overflow_msg):
-        ops.calculate(operator.add, 1, int_max)
+    def test_raises_for_too_small_result(self, int_min: int, overflow_msg: str):
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_int(operator.sub, int_min, 1)
 
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_int(operator.sub, 1, int_min)
 
-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)
+class TestCalcIntFloat:
+    @pytest.mark.parametrize(
+        "op,lval,rval,expected",
+        (
+            (operator.add, 1, 1.0, 2),
+            (operator.sub, 2, 1.0, 1),
+            (operator.mul, 1, 2.0, 2),
+            (operator.truediv, 1, 0.5, 2),
+        ),
+        ids=("+", "-", "*", "/"),
+    )
+    def test_arithmetic_ops(self, op, lval: int, rval: float, expected: int):
+        result = ops.calc_int_float(op, lval, rval)
+
+        assert result == expected
+        assert isinstance(result, int)
+
+    def test_raises_for_too_large_arg(
+        self,
+        int_max: int,
+        float_max: float,
+        overflow_msg: str,
+    ):
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_float(operator.add, int_max + 1, 1)
+
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_float(operator.add, 1, float_max + 1)
+
+    def test_raises_for_too_small_arg(
+        self,
+        int_min: int,
+        float_min: float,
+        overflow_msg: str,
+    ):
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_float(operator.add, int_min - 1, 1)
+
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_float(operator.add, 1, float_min - 1)
+
+    def test_raises_for_too_large_result(
+        self,
+        int_max: int,
+        float_max: float,
+        overflow_msg: str,
+    ):
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_float(operator.add, int_max, 1)
+
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_float(operator.add, 1, float_max)
+
+    @pytest.mark.parametrize(
+        "value",
+        (
+            pytest.param(
+                1024,
+                marks=pytest.mark.xfail(
+                    reason="TBD",
+                    raises=pytest.fail.Exception,
+                    strict=True,
+                ),
+            ),
+            1024.1,
+        ),
+    )
+    def test_raises_for_most_too_small_results(
+        self,
+        value: float,
+        int_min: int,
+        overflow_msg: str,
+    ):
+        with pytest.raises(OverflowError, match=overflow_msg):
+            ops.calc_int_float(operator.sub, int_min, value)