Add support for operator overloading of Duration (#1387)

* Add support for operator overloading of `Duration`

This PR adds support for operator overloading of the `Duration` class in
python. This should improve the overall UX of the rclpy library.
---------

Signed-off-by: Arjo Chakravarty <[email protected]>

rclpy/rclpy/duration.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import math
 from typing import Union
 
 import builtin_interfaces.msg
@@ -49,6 +49,28 @@ def __str__(self) -> str:
             return 'Infinite'
         return f'{self.nanoseconds} nanoseconds'
 
+    def __add__(self, other: 'Duration') -> 'Duration':
+        if isinstance(other, Duration):
+            return Duration(nanoseconds=other.nanoseconds + self.nanoseconds)
+        return NotImplemented
+
+    def __sub__(self, other: 'Duration') -> 'Duration':
+        if isinstance(other, Duration):
+            return Duration(nanoseconds=self.nanoseconds - other.nanoseconds)
+        return NotImplemented
+
+    def __mul__(self, other: Union[int, float]) -> 'Duration':
+        if isinstance(other, int):
+            return Duration(nanoseconds=self.nanoseconds * other)
+        if isinstance(other, float):
+            if not math.isfinite(other):
+                if other == float('inf'):
+                    return Infinite
+                else:
+                    raise ValueError("Can't multiply duration with nan")
+            return Duration(nanoseconds=int(self.nanoseconds * other))
+        return NotImplemented
+
     def __eq__(self, other: object) -> bool:
         if isinstance(other, Duration):
             return self.nanoseconds == other.nanoseconds

rclpy/test/test_time.py

@@ -15,6 +15,7 @@
 import unittest
 
 from rclpy.clock_type import ClockType
+from rclpy.constants import S_TO_NS
 from rclpy.duration import Duration
 from rclpy.duration import Infinite
 from rclpy.time import Time
@@ -68,6 +69,17 @@ def test_duration_construction(self) -> None:
         with self.assertRaises(OverflowError):
             Duration(nanoseconds=-2**63 - 1)
 
+    def test_duration_operators(self) -> None:
+        duration1 = Duration(seconds=1, nanoseconds=0)
+        duration2 = Duration(seconds=2, nanoseconds=0)
+        assert (duration1 + duration2).nanoseconds == 3 * S_TO_NS
+        assert (duration2 - duration1).nanoseconds == 1 * S_TO_NS
+        assert (duration1 * 2) == duration2
+        assert (duration2 * 0.5) == duration1
+        assert (duration2 * float('inf')) == Infinite
+        with self.assertRaises(ValueError):
+            duration2 * float('NaN')
+
     def test_time_operators(self) -> None:
         time1 = Time(nanoseconds=1, clock_type=ClockType.STEADY_TIME)