feat(processing_time_visualizer): enable to wait until topic published

common/autoware_debug_tools/autoware_debug_tools/processing_time_visualizer/node.py

@@ -1,3 +1,4 @@
+import argparse
 import curses
 import json
 import time
@@ -18,8 +19,10 @@
 
 
 class ProcessingTimeVisualizer(Node):
-    def __init__(self):
+    def __init__(self, topic_waiting_sec=1.0, topic_name=None):
         super().__init__("processing_time_visualizer" + str(uuid.uuid4()).replace("-", "_"))
+        self.topic_name = topic_name
+        self.topic_waiting_sec = topic_waiting_sec
         self.subscriber = self.subscribe_processing_time_tree()
         self.quit_option = None
         self.trees: Dict[str, ProcessingTimeTree] = {}
@@ -32,35 +35,65 @@ def __init__(self):
 
         self.create_timer(0.1, self.update_screen)
 
-    def subscribe_processing_time_tree(self):
-        topics = []
+    def search_for_topic(self, topic_name: str) -> bool:
+        for name, topic_types in self.get_topic_names_and_types():
+            if name == topic_name:
+                return True
+        return False
 
+    def wait_for_topics(self, condition_func, wait_sec: float) -> bool:
         s = time.time()
-        while True:
-            for topic_name, topic_types in self.get_topic_names_and_types():
-                for topic_type in topic_types:
-                    if (
-                        topic_type == "tier4_debug_msgs/msg/ProcessingTimeTree"
-                        and topic_name not in topics
-                    ):
-                        topics.append(topic_name)
-
-            if time.time() - s > 1.0:
-                break
-
-        if len(topics) == 0:
-            self.get_logger().info("No ProcessingTimeTree topic found")
-            self.get_logger().info("Exiting...")
-            exit(1)
-        else:
-            self.topic_name = curses.wrapper(select_topic, topics)
-            subscriber = self.create_subscription(
-                ProcessingTimeTreeMsg,
-                self.topic_name,
-                self.callback,
-                10,
+        while time.time() - s < wait_sec:
+            if condition_func():
+                return True
+        return False
+
+    def subscribe_processing_time_tree(self):
+        # if topic name is specified with -t option
+        if self.topic_name:
+            topic_found = self.wait_for_topics(
+                lambda: self.search_for_topic(self.topic_name), self.topic_waiting_sec
             )
 
+            if not topic_found:
+                self.get_logger().info(f"Specified topic '{self.topic_name}' not found.")
+                self.get_logger().info("Exiting...")
+                exit(1)
+            else:
+                subscriber = self.create_subscription(
+                    ProcessingTimeTreeMsg,
+                    self.topic_name,
+                    self.callback,
+                    10,
+                )
+        else:  # if topic name is not specified
+            topics = []
+
+            def condition_func():
+                for name, topic_types in self.get_topic_names_and_types():
+                    for topic_type in topic_types:
+                        if (
+                            topic_type == "tier4_debug_msgs/msg/ProcessingTimeTree"
+                            and name not in topics
+                        ):
+                            topics.append(name)
+                return len(topics) > 0
+
+            topic_found = self.wait_for_topics(condition_func, self.topic_waiting_sec)
+
+            if not topic_found:
+                self.get_logger().info("No ProcessingTimeTree topic found")
+                self.get_logger().info("Exiting...")
+                exit(1)
+            else:
+                self.topic_name = curses.wrapper(select_topic, topics)
+                subscriber = self.create_subscription(
+                    ProcessingTimeTreeMsg,
+                    self.topic_name,
+                    self.callback,
+                    10,
+                )
+
         return subscriber
 
     def update_screen(self):
@@ -107,9 +140,18 @@ def callback(self, msg: ProcessingTimeTreeMsg):
 
 
 def main(args=None):
+    parser = argparse.ArgumentParser(description="Processing Time Visualizer")
+    parser.add_argument("-t", "--topic", type=str, help="Specify the topic name to subscribe to")
+    parser.add_argument(
+        "-w", "--waiting-sec", type=float, default=1.0, help="Waiting time for topic"
+    )
+    parsed_args = parser.parse_args(args)
+
     rclpy.init(args=args)
     try:
-        node = ProcessingTimeVisualizer()
+        node = ProcessingTimeVisualizer(
+            topic_name=parsed_args.topic, topic_waiting_sec=parsed_args.waiting_sec
+        )
     except KeyboardInterrupt:
         exit_curses()
         return