Rework TFListenerWrapper API

Signed-off-by: Michel Hidalgo <[email protected]>

bdai_ros2_wrappers/bdai_ros2_wrappers/action_handle.py

@@ -17,7 +17,7 @@ class ActionHandle(object):
     as holding the various callbacks for sending an ActionGoal (cancel, failure, feedback, result)
     """
 
-    def __init__(self, action_name: str, logger: Optional[RcutilsLogger] = None, *, context: Optional[Context] = None):
+    def __init__(self, action_name: str, logger: Optional[RcutilsLogger] = None, context: Optional[Context] = None):
         """Constructor
 
         Args:

bdai_ros2_wrappers/bdai_ros2_wrappers/tf_listener_wrapper.py

@@ -1,116 +1,60 @@
 # Copyright (c) 2023 Boston Dynamics AI Institute Inc.  All rights reserved.
-import time
-from typing import Optional, Tuple
 
-import rclpy
+from typing import Optional
+
 from geometry_msgs.msg import TransformStamped
 from rclpy.node import Node
+from rclpy.task import Future
 from rclpy.time import Duration, Time
-from tf2_ros import ExtrapolationException, TransformException
 from tf2_ros.buffer import Buffer
 from tf2_ros.transform_listener import TransformListener
 
 import bdai_ros2_wrappers.scope as scope
+from bdai_ros2_wrappers.futures import wait_for_future
 
 
-class TFListenerWrapper(object):
-    def __init__(
-        self,
-        node: Optional[Node] = None,
-        wait_for_transform: Optional[Tuple[str, str]] = None,
-        cache_time_s: Optional[int] = None,
-    ) -> None:
+class TFListenerWrapper:
+    def __init__(self, node: Optional[Node] = None, cache_time_s: Optional[float] = None) -> None:
         node = node or scope.node()
         if node is None:
             raise ValueError("no ROS 2 node available (did you use bdai_ros2_wrapper.process.main?)")
         self._node = node
-
+        cache_time_py = None
         if cache_time_s is not None:
-            cache_time_py = Duration(seconds=cache_time_s)
-        else:
-            cache_time_py = None
+            cache_time_py = Duration(nanoseconds=cache_time_s * 1e9)
         self._tf_buffer = Buffer(cache_time=cache_time_py)
-        self._tf_listener = TransformListener(self._tf_buffer, self._node, spin_thread=False)
-
-        if wait_for_transform is not None and len(wait_for_transform) == 2:
-            self.wait_for_init(wait_for_transform[0], wait_for_transform[1])
+        self._tf_listener = TransformListener(self._tf_buffer, self._node)
 
     @property
     def buffer(self) -> Buffer:
         return self._tf_buffer
 
-    def wait_for_init(self, from_frame: str, to_frame: str) -> None:
-        """
-        Waits for transform from from_frame to to_frame to become available and prints initializing statements.
-        """
-        # Wait for the buffer to fill to avoid annoying warnings.
-        self._node.get_logger().info(
-            "Waiting for TF to contain transform from " + str(from_frame) + " to " + str(to_frame)
-        )
-        self.wait_for_transform(from_frame, to_frame)
-        self._node.get_logger().info("TF initialized")
+    def shutdown(self) -> None:
+        self._tf_listener.unregister()
 
-    def wait_for_transform(self, from_frame: str, to_frame: str) -> None:
-        """
-        Wait for the transform from from_frame to to_frame to become available.
-        """
-        # We need a while instead of the usual wait because we use a different node and a thread for our listener
-        while rclpy.ok(context=self._node.context):
-            try:
-                self.lookup_a_tform_b(from_frame, to_frame)
-                break
-            except TransformException:
-                time.sleep(0.1)
+    def wait_for_a_tform_b_async(self, frame_a: str, frame_b: str, time: Optional[Time] = None) -> Future:
+        if time is None:
+            time = Time()
+        return self._tf_buffer.wait_for_transform_async(frame_a, frame_b, time)
 
-    def _internal_lookup_a_tform_b(
-        self,
-        frame_a: str,
-        frame_b: str,
-        transform_time: Optional[float] = None,
-        timeout: Optional[float] = None,
-        wait_for_frames: bool = False,
-    ) -> TransformStamped:
+    def wait_for_a_tform_b(
+        self, frame_a: str, frame_b: str, time: Optional[Time] = None, timeout_sec: Optional[float] = None
+    ) -> bool:
         """
-        Parameters:
-            frame_a: Base frame for transform.  The transform returned will be frame_a_t_frame_b
-            frame_b: Tip frame for transform.  The transform returned will be frame_a_t_frame_b
-            transform_time: The time at which to look up the transform.  If left at None, will be the most recent
-                transform available
-            timeout: The time to wait for the transform to become available if the transform_time is beyond the most
-                recent transform in the buffer.
-            wait_for_frames: If true, waits timeout amount of time for a path to exist from frame_a to frame_b in the
-                buffer.  If false, this will return immediately if a path does not exist even if timeout is not
-                None.  Note that wait_for_transform can be used to wait indefinitely for a transform to become
-                available.
-        Returns:
-            The transform frame_a_t_frame_b at the time specified.
-        Raises:
-            All the possible TransformExceptions.
+        Wait for the transform from from_frame to to_frame to become available.
         """
-        if transform_time is None:
-            transform_time = Time()
-        if timeout is None or not wait_for_frames:
-            timeout_py = Duration()
-        else:
-            timeout_py = Duration(seconds=timeout)
-        start_time = time.time()
-        while rclpy.ok(context=self._node.context):
-            try:
-                return self._tf_buffer.lookup_transform(frame_a, frame_b, time=transform_time, timeout=timeout_py)
-            except ExtrapolationException as e:
-                if "future" not in str(e) or timeout is None:
-                    raise e  # Waiting won't help with this
-                now = time.time()
-                if now - start_time > timeout:
-                    raise e
-                time.sleep(0.01)
+        future = self.wait_for_a_tform_b_async(frame_a, frame_b, time)
+        if not wait_for_future(future, timeout_sec, context=self._node.context):
+            future.cancel()
+            return False
+        return True
 
     def lookup_a_tform_b(
         self,
         frame_a: str,
         frame_b: str,
-        transform_time: Optional[float] = None,
-        timeout: Optional[float] = None,
+        transform_time: Optional[Time] = None,
+        timeout_sec: Optional[float] = None,
         wait_for_frames: bool = False,
     ) -> TransformStamped:
         """
@@ -119,36 +63,40 @@ def lookup_a_tform_b(
             frame_b: Tip frame for transform.  The transform returned will be frame_a_t_frame_b
             transform_time: The time at which to look up the transform.  If left at None, will be the most recent
                 transform available
-            timeout: The time to wait for the transform to become available if the transform_time is beyond the most
-                recent transform in the buffer.
+            timeout_sec: The time to wait for the transform to become available if the requested time is beyond
+                the most recent transform in the buffer. If set to 0, it will not wait. If left at None, it will
+                wait indefinitely.
             wait_for_frames: If true, waits timeout amount of time for a path to exist from frame_a to frame_b in the
-                buffer.  If false, this will return immediately if a path does not exist even if timeout is not
-                None.  Note that wait_for_transform can be used to wait indefinitely for a transform to become
-                available.
+                buffer. If false, this will return immediately if a path does not exist even if timeout is not
+                None. Note that wait_for_a_tform_b can also be used to wait for a transform to become available.
         Returns:
             The transform frame_a_t_frame_b at the time specified.
         Raises:
             All the possible TransformExceptions.
         """
-        return self._internal_lookup_a_tform_b(frame_a, frame_b, transform_time, timeout, wait_for_frames)
+        if not wait_for_frames:
+            timeout_sec = 0.0
+        if timeout_sec != 0.0:
+            self.wait_for_a_tform_b(frame_a, frame_b, transform_time, timeout_sec)
+        return self._tf_buffer.lookup_transform(frame_a, frame_b, transform_time)
 
     def lookup_latest_timestamp(
-        self, frame_a: str, frame_b: str, timeout: Optional[float] = None, wait_for_frames: bool = False
+        self, frame_a: str, frame_b: str, timeout_sec: Optional[float] = None, wait_for_frames: bool = False
     ) -> Time:
         """
         Parameters:
             frame_a: Base frame for transform.  The transform returned will be frame_a_t_frame_b
             frame_b: Tip frame for transform.  The transform returned will be frame_a_t_frame_b
-            timeout: The time to wait for the transform to become available if the transform_time is beyond the most
-                recent transform in the buffer.
+            timeout_sec: The time to wait for the transform to become available if the requested time is beyond
+                the most recent transform in the buffer. If set to 0, it will not wait. If left at None, it will
+                wait indefinitely.
             wait_for_frames: If true, waits timeout amount of time for a path to exist from frame_a to frame_b in the
-                buffer.  If false, this will return immediately if a path does not exist even if timeout is not
-                None.  Note that wait_for_transform can be used to wait indefinitely for a transform to become
-                available.
+                buffer. If false, this will return immediately if a path does not exist even if timeout is not
+                None. Note that wait_for_a_tform_b can be used to wait for a transform to become available.
         Returns:
             The timestamp from the latest recorded transform frame_a_t_frame_b
         Raises:
             All the possible TransformExceptions.
         """
-        transform = self._internal_lookup_a_tform_b(frame_a, frame_b, None, timeout, wait_for_frames)
+        transform = self.lookup_a_tform_b(frame_a, frame_b, timeout_sec, wait_for_frames)
         return Time.from_msg(transform.header.stamp)

bdai_ros2_wrappers/test/test_tf_listener_wrapper.py

@@ -9,7 +9,7 @@
 from tf2_ros import ExtrapolationException, LookupException, TransformBroadcaster
 
 from bdai_ros2_wrappers.node import Node
-from bdai_ros2_wrappers.process import ROSAwareScope
+from bdai_ros2_wrappers.scope import ROSAwareScope
 from bdai_ros2_wrappers.tf_listener_wrapper import TFListenerWrapper
 
 ROBOT = "test_robot"
@@ -64,7 +64,7 @@ def test_non_existant_transform(ros: ROSAwareScope, tf_pair: Tuple[MockTfPublish
     assert ros.node is not None
     timestamp = ros.node.get_clock().now()
     with pytest.raises(LookupException):
-        tf_listener.lookup_a_tform_b(FRAME_ID, CHILD_FRAME_ID, timestamp)
+        tf_listener.lookup_a_tform_b(FRAME_ID, CHILD_FRAME_ID, timestamp, timeout_sec=0)
 
 
 def test_non_existant_transform_timeout(
@@ -75,8 +75,9 @@ def test_non_existant_transform_timeout(
     timestamp = ros.node.get_clock().now()
     start = time.time()
     with pytest.raises(LookupException):
-        tf_listener.lookup_a_tform_b(FRAME_ID, CHILD_FRAME_ID, timestamp, timeout=20.0)
-    assert time.time() - start < 10.0
+        tf_listener.lookup_a_tform_b(FRAME_ID, CHILD_FRAME_ID, timestamp, timeout_sec=2.0, wait_for_frames=True)
+    elapsed_time = time.time() - start
+    assert elapsed_time >= 2.0
 
 
 def test_existing_transform(ros: ROSAwareScope, tf_pair: Tuple[MockTfPublisherNode, TFListenerWrapper]) -> None:
@@ -127,7 +128,7 @@ def delayed_publish() -> None:
     time.sleep(0.2)
     timestamp = ros.node.get_clock().now() + Duration(seconds=delay)
     with pytest.raises(ExtrapolationException):
-        tf_listener.lookup_a_tform_b(FRAME_ID, CHILD_FRAME_ID, timestamp, timeout=0.5)
+        tf_listener.lookup_a_tform_b(FRAME_ID, CHILD_FRAME_ID, timestamp, timeout_sec=0.5, wait_for_frames=True)
 
 
 def test_future_transform_wait(ros: ROSAwareScope, tf_pair: Tuple[MockTfPublisherNode, TFListenerWrapper]) -> None:
@@ -148,5 +149,5 @@ def delayed_publish() -> None:
     ros.executor.create_task(delayed_publish)
 
     timestamp += Duration(seconds=delay)
-    t = tf_listener.lookup_a_tform_b(FRAME_ID, CHILD_FRAME_ID, timestamp, wait_for_frames=True, timeout=2)
+    t = tf_listener.lookup_a_tform_b(FRAME_ID, CHILD_FRAME_ID, timestamp, timeout_sec=2.0, wait_for_frames=True)
     assert equal_transform(t.transform, trans)