RobotecAI · maciejmajek · Apr 17, 2025 · Apr 17, 2025 · Apr 17, 2025 · Apr 17, 2025
diff --git a/src/rai_core/rai/communication/base_connector.py b/src/rai_core/rai/communication/base_connector.py
@@ -22,7 +22,9 @@
     Dict,
     Generic,
     Optional,
+    Type,
     TypeVar,
+    get_args,
 )
 from uuid import uuid4
 
@@ -68,6 +70,15 @@ def __init__(self, callback_max_workers: int = 4):
             max_workers=self.callback_max_workers
         )
 
+        if not hasattr(self, "__orig_bases__"):
+            self.__orig_bases__ = {}
+            raise ConnectorException(
+                f"Error while instantiating {str(self.__class__)}: "
+                "Message type T derived from BaseMessage needs to be provided"
+                " e.g. Connector[MessageType]()"
+            )
+        self.T_class: Type[T] = get_args(self.__orig_bases__[-1])[0]
+
     def _generate_handle(self) -> str:
         return str(uuid4())
 

diff --git a/src/rai_core/rai/communication/ros2/connectors/__init__.py b/src/rai_core/rai/communication/ros2/connectors/__init__.py
@@ -13,12 +13,14 @@
 # limitations under the License.
 
 from .action_mixin import ROS2ActionMixin
-from .connector import ROS2Connector
+from .base import ROS2BaseConnector
 from .hri_connector import ROS2HRIConnector
+from .ros2_connector import ROS2Connector
 from .service_mixin import ROS2ServiceMixin
 
 __all__ = [
     "ROS2ActionMixin",
+    "ROS2BaseConnector",
     "ROS2Connector",
     "ROS2HRIConnector",
     "ROS2ServiceMixin",

diff --git a/...ommunication/ros2/connectors/connector.py → ...rai/communication/ros2/connectors/base.py b/...ommunication/ros2/connectors/connector.py → ...rai/communication/ros2/connectors/base.py
@@ -16,7 +16,7 @@
 import time
 import uuid
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar
 
 import rclpy
 import rclpy.executors
@@ -38,8 +38,10 @@
 from rai.communication.ros2.connectors.service_mixin import ROS2ServiceMixin
 from rai.communication.ros2.messages import ROS2Message
 
+T = TypeVar("T", bound=ROS2Message)
 
-class ROS2Connector(ROS2ActionMixin, ROS2ServiceMixin, BaseConnector[ROS2Message]):
+
+class ROS2BaseConnector(ROS2ActionMixin, ROS2ServiceMixin, BaseConnector[T]):
     """ROS2-specific implementation of the ARIConnector.
 
        This connector provides functionality for ROS2 communication through topics,
@@ -110,9 +112,9 @@ def __init__(
         self._thread.start()
 
         # cache for last received messages
-        self.last_msg: Dict[str, ROS2Message] = {}
+        self.last_msg: Dict[str, T] = {}
 
-    def last_message_callback(self, source: str, msg: ROS2Message):
+    def last_message_callback(self, source: str, msg: T):
         self.last_msg[source] = msg
 
     def get_topics_names_and_types(self) -> List[Tuple[str, List[str]]]:
@@ -126,7 +128,7 @@ def get_actions_names_and_types(self) -> List[Tuple[str, List[str]]]:
 
     def send_message(
         self,
-        message: ROS2Message,
+        message: T,
         target: str,
         *,
         msg_type: str,  # TODO: allow msg_type to be None, add auto topic type detection
@@ -143,19 +145,19 @@ def send_message(
         )
 
     def general_callback_preprocessor(self, message: Any):
-        return ROS2Message(payload=message, metadata={"msg_type": str(type(message))})
+        return self.T_class(payload=message, metadata={"msg_type": str(type(message))})
 
     def register_callback(
         self,
         source: str,
-        callback: Callable[[ROS2Message | Any], None],
+        callback: Callable[[T | Any], None],
         raw: bool = False,
         *,
         msg_type: Optional[str] = None,
         qos_profile: Optional[QoSProfile] = None,
         auto_qos_matching: bool = True,
         **kwargs: Any,
-    ):
+    ) -> str:
         exists = self._topic_api.subscriber_exists(source)
         if not exists:
             self._topic_api.create_subscriber(
@@ -165,7 +167,7 @@ def register_callback(
                 qos_profile=qos_profile,
                 auto_qos_matching=auto_qos_matching,
             )
-        super().register_callback(source, callback, raw=raw)
+        return super().register_callback(source, callback, raw=raw)
 
     def receive_message(
         self,
@@ -176,7 +178,7 @@ def receive_message(
         qos_profile: Optional[QoSProfile] = None,
         auto_qos_matching: bool = True,
         **kwargs: Any,
-    ) -> ROS2Message:
+    ) -> T:
         if self._topic_api.subscriber_exists(source):
             # trying to hit cache first
             if source in self.last_msg:

diff --git a/src/rai_core/rai/communication/ros2/connectors/hri_connector.py b/src/rai_core/rai/communication/ros2/connectors/hri_connector.py
@@ -13,22 +13,14 @@
 # limitations under the License.
 
 import logging
-import threading
 import uuid
 from typing import Any, Callable, List, Literal, Optional, Tuple, Union
 
-from rclpy.executors import MultiThreadedExecutor
-from rclpy.node import Node
-
-from rai.communication import HRIConnector
 from rai.communication.ros2.api import (
     ConfigurableROS2TopicAPI,
-    ROS2ActionAPI,
-    ROS2ServiceAPI,
     TopicConfig,
 )
-from rai.communication.ros2.connectors.action_mixin import ROS2ActionMixin
-from rai.communication.ros2.connectors.service_mixin import ROS2ServiceMixin
+from rai.communication.ros2.connectors.base import ROS2BaseConnector
 from rai.communication.ros2.messages import ROS2HRIMessage
 
 try:
@@ -37,7 +29,7 @@
     logging.warning("rai_interfaces is not installed, ROS 2 HRIMessage will not work.")
 
 
-class ROS2HRIConnector(ROS2ActionMixin, ROS2ServiceMixin, HRIConnector[ROS2HRIMessage]):
+class ROS2HRIConnector(ROS2BaseConnector[ROS2HRIMessage]):
     def __init__(
         self,
         node_name: str = f"rai_ros2_hri_connector_{str(uuid.uuid4())[-12:]}",
@@ -50,6 +42,8 @@ def __init__(
         configured_sources = [
             source[0] if isinstance(source, tuple) else source for source in sources
         ]
+        self.configured_targets = configured_targets
+        self.configured_sources = configured_sources
 
         _targets = [
             (
@@ -67,22 +61,11 @@ def __init__(
             )
             for source in sources
         ]
-
-        self._node = Node(node_name)
+        super().__init__(node_name=node_name)
         self._topic_api = ConfigurableROS2TopicAPI(self._node)
-        self._service_api = ROS2ServiceAPI(self._node)
-        self._actions_api = ROS2ActionAPI(self._node)
-
         self._configure_publishers(_targets)
         self._configure_subscribers(_sources)
 
-        super().__init__(configured_targets, configured_sources)
-
-        self._executor = MultiThreadedExecutor()
-        self._executor.add_node(self._node)
-        self._thread = threading.Thread(target=self._executor.spin)
-        self._thread.start()
-
     def _configure_publishers(self, targets: List[Tuple[str, TopicConfig]]):
         for target in targets:
             self._topic_api.configure_publisher(target[0], target[1])

diff --git a/src/rai_core/rai/communication/ros2/connectors/ros2_connector.py b/src/rai_core/rai/communication/ros2/connectors/ros2_connector.py
@@ -0,0 +1,20 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+from rai.communication.ros2.connectors.base import ROS2BaseConnector
+from rai.communication.ros2.messages import ROS2Message
+
+
+class ROS2Connector(ROS2BaseConnector[ROS2Message]):
+    pass
diff --git a/src/rai_core/rai/communication/ros2/messages.py b/src/rai_core/rai/communication/ros2/messages.py
@@ -41,7 +41,7 @@ class ROS2Message(BaseMessage):
     pass
 
 
-class ROS2HRIMessage(HRIMessage):
+class ROS2HRIMessage(HRIMessage, ROS2Message):
     @classmethod
     def from_ros2(
         cls,

diff --git a/tests/rai_sim/test_o3de_bridge.py b/tests/rai_sim/test_o3de_bridge.py
@@ -301,20 +301,6 @@ def test_send_message_signature(self):
             f"Parameter names do not match, expected: {list(expected_params.keys())}, got: {list(parameters.keys())}",
         )
 
-        for param_name, expected_type in expected_params.items():
-            param = parameters[param_name]
-            self.assertEqual(
-                self.resolve_annotation(param.annotation),
-                self.resolve_annotation(expected_type),
-                f"Parameter '{param_name}' has incorrect type, expected: {expected_type}, got: {param.annotation}",
-            )
-
-        self.assertIs(
-            signature.return_annotation,
-            inspect.Signature.empty,
-            "send_message should have no return value",
-        )
-
     def test_receive_message_signature(self):
         signature = inspect.signature(self.connector.receive_message)
         parameters = signature.parameters
@@ -341,12 +327,6 @@ def test_receive_message_signature(self):
                 f"Parameter '{param_name}' has incorrect type, expected: {expected_type}, got: {param.annotation}",
             )
 
-        self.assertIs(
-            signature.return_annotation,
-            ROS2Message,
-            f"Return type is incorrect, expected: ROS2Message, got: {signature.return_annotation}",
-        )
-
     def test_get_topics_names_and_types_signature(self):
         signature = inspect.signature(self.connector.get_topics_names_and_types)
         parameters = signature.parameters