-Add '+' and '-' annotations for Timestamp and Duration in Python

-Add support for "Duration - Duration" and "Duration - timedelata" in Python

PiperOrigin-RevId: 747487973

Author: anandolee

python/google/protobuf/internal/duration_test.py

@@ -65,6 +65,32 @@ def test_duration_construction(self):
     )
     self.assertEqual(expected_td, message.optional_duration.ToTimedelta())
 
+  def test_duration_sub_annotation(self):
+    dt = datetime.datetime.now()
+    dr = duration_pb2.Duration()
+    td = datetime.timedelta(microseconds=123)
+    # datetime - Duration
+    self.assertEqual(dt - dr, dt - duration.to_timedelta(dr))
+    # timedelta - Duration and Duration - Duration
+    self.assertEqual(td - dr, duration.from_timedelta(td) - dr)
+    # Duration - timedelta
+    self.assertEqual(dr - td, dr - duration.from_timedelta(td))
+
+  def test_duration_add_annotation(self):
+    dt = datetime.datetime.now()
+    dr = duration_pb2.Duration()
+    dr2 = duration_pb2.Duration(seconds=100)
+    # datetime + Duration and Duration + datetime
+    self.assertEqual(dt + dr, dr + dt)
+    message = well_known_types_test_pb2.WKTMessage(optional_timestamp=dt)
+    # Duration + Timestamp
+    self.assertEqual(dr + message.optional_timestamp, dr + dt)
+    td = datetime.timedelta(microseconds=123)
+    # Duration + timedelta and timedelta + Duration
+    self.assertEqual(dr + td, td + dr)
+    # Duration + Duration
+    self.assertEqual(dr + dr2, dr2 + dr)
+
 
 if __name__ == '__main__':
   unittest.main()

python/google/protobuf/internal/timestamp_test.py

@@ -64,6 +64,30 @@ def test_timstamp_construction(self):
         optional_timestamp=datetime.datetime.today()
     )
 
+  def test_timestamp_sub_annotation(self):
+    t1 = timestamp_pb2.Timestamp()
+    t2 = timestamp_pb2.Timestamp()
+    dt = datetime.datetime.now()
+    td = datetime.timedelta(hours=0)
+    msg = well_known_types_test_pb2.WKTMessage(optional_duration=td)
+    # Timestamp - datetime
+    self.assertEqual(t1 - dt, t2 - dt)
+    # Timestamp - Timestamp
+    self.assertEqual(t1 - t2, t2 - t1)
+    # datetime - Timestamp
+    self.assertEqual(dt - t1, dt - t2)
+    # Timestamp - timedelta and Timestamp - Duration
+    self.assertEqual(t1 - td, t2 - msg.optional_duration)
+
+  def test_timestamp_add_annotation(self):
+    ts = timestamp_pb2.Timestamp()
+    td = datetime.timedelta(hours=0)
+    msg = well_known_types_test_pb2.WKTMessage(optional_duration=td)
+    # Timestamp + timedelta and timedelta + Timestamp
+    self.assertEqual(ts + td, td + ts)
+    # Timestamp + Duration and Duration + Timestamp
+    self.assertEqual(ts + msg.optional_duration, msg.optional_duration + ts)
+
 
 if __name__ == '__main__':
   unittest.main()

python/google/protobuf/internal/well_known_types.py

@@ -475,8 +475,11 @@ def __add__(self, value) -> Union[datetime.datetime, datetime.timedelta]:
 
   __radd__ = __add__
 
-  def __rsub__(self, dt) -> Union[datetime.datetime, datetime.timedelta]:
-    return dt - self.ToTimedelta()
+  def __sub__(self, value) -> datetime.timedelta:
+    return self.ToTimedelta() - value
+
+  def __rsub__(self, value) -> Union[datetime.datetime, datetime.timedelta]:
+    return value - self.ToTimedelta()
 
 
 def _CheckDurationValid(seconds, nanos):