Support datetime/timedelta arithmetic between csp edges (#676)

Signed-off-by: Adam Glustein <adam.glustein@point72.com>
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: AdamGlustein

csp/math.py

@@ -1,4 +1,5 @@
 import math
+from datetime import datetime, timedelta
 from functools import lru_cache
 from typing import List, TypeVar, get_origin
 
@@ -270,6 +271,18 @@ def generic_type(x: ts["T"], y: ts["T"]) -> ts[generic_out_type]:
         if csp.valid(x, y):
             return op_lambda(x, y)
 
+    # Special case: datetime - datetime returns timedelta
+    @_node_internal_use(name=name)
+    def datetime_sub_type(x: ts[datetime], y: ts[datetime]) -> ts[timedelta]:
+        if csp.valid(x, y):
+            return op_lambda(x, y)
+
+    # Special case: datetime +/- timedelta returns datetime
+    @_node_internal_use(name=name)
+    def datetime_timedelta_type(x: ts[datetime], y: ts[timedelta]) -> ts[datetime]:
+        if csp.valid(x, y):
+            return op_lambda(x, y)
+
     def comp(x: ts["T"], y: ts["U"]):
         if get_origin(x.tstype.typ) in [Numpy1DArray, NumpyNDArray] or get_origin(y.tstype.typ) in [
             Numpy1DArray,
@@ -280,6 +293,10 @@ def comp(x: ts["T"], y: ts["U"]):
             return float_type(x, y)
         elif x.tstype.typ is int and y.tstype.typ is int:
             return int_type(x, y)
+        elif name == "sub" and x.tstype.typ is datetime and y.tstype.typ is datetime:
+            return datetime_sub_type(x, y)
+        elif name in ("add", "sub") and x.tstype.typ is datetime and y.tstype.typ is timedelta:
+            return datetime_timedelta_type(x, y)
         return generic_type(x, y)
 
     comp.__name__ = name

csp/tests/test_math.py

@@ -281,6 +281,33 @@ def graph(use_promotion: bool):
                     [v[1] for v in results[op.__name__ + "-rev"]], [comp(y, x) for x, y in zip(xv, yv)], op.__name__
                 )
 
+    def test_arithmetic_time_ops(self):
+        """Test datetime/timedelta arithmetic operations between edges."""
+
+        @csp.graph
+        def graph():
+            dt1 = datetime(2020, 1, 1, 12, 0, 0)
+            dt2 = datetime(2020, 1, 1, 10, 0, 0)
+            td = timedelta(hours=2)
+            dt1_edge = csp.const(dt1)
+            dt2_edge = csp.const(dt2)
+            td_edge = csp.const(td)
+
+            # datetime - datetime -> timedelta
+            csp.add_graph_output("dt_sub_dt", dt1_edge - dt2_edge)
+
+            # datetime + timedelta -> datetime
+            csp.add_graph_output("dt_plus_td", dt1_edge + td_edge)
+
+            # datetime - timedelta -> datetime
+            csp.add_graph_output("dt_minus_td", dt1_edge - td_edge)
+
+        st = datetime(2020, 1, 1)
+        results = csp.run(graph, starttime=st, endtime=st + timedelta(seconds=1))
+        self.assertEqual(results["dt_sub_dt"][0][1], timedelta(hours=2))
+        self.assertEqual(results["dt_plus_td"][0][1], datetime(2020, 1, 1, 14, 0, 0))
+        self.assertEqual(results["dt_minus_td"][0][1], datetime(2020, 1, 1, 10, 0, 0))
+
     def test_boolean_ops(self):
         def graph():
             x = csp.default(csp.curve(bool, [(timedelta(seconds=s), s % 2 == 0) for s in range(1, 20)]), False)