feat: add ROS2 message support and enhance publishing functionality

Implement functions to register, publish, and subscribe to ROS2 messages, improving interoperability with ROS2 systems. Adjust type support to handle both protobuf and ROS2 message types, streamlining the process of message serialization and ensuring robustness in message handling.
AimRT · AimRT · Nov 21, 2024 · Nov 7, 2024 · Nov 7, 2024 · Nov 8, 2024
commit 455a2fcc19c0f3f8bfd02d370f8b2ec6371042fd
diff --git a/src/runtime/python_runtime/CMakeLists.txt b/src/runtime/python_runtime/CMakeLists.txt
@@ -26,7 +26,8 @@ target_sources(${CUR_TARGET_NAME} PRIVATE ${src})
 # Set link libraries of target
 target_link_libraries(
   ${CUR_TARGET_NAME}
-  PRIVATE $<LINK_LIBRARY:WHOLE_ARCHIVE,aimrt::runtime::core>)
+  PRIVATE $<LINK_LIBRARY:WHOLE_ARCHIVE,aimrt::runtime::core>
+          aimrt::interface::aimrt_module_ros2_interface)
 
 # Set misc of target
 set_target_properties(${CUR_TARGET_NAME} PROPERTIES OUTPUT_NAME "aimrt_python_runtime")

diff --git a/src/runtime/python_runtime/__init__.py b/src/runtime/python_runtime/__init__.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2023, AgiBot Inc.
 # All rights reserved.
 
+from .aimrt_py_chn import (Publish, PublishRos2Message, RegisterPublishType,
+                           Subscribe, SubscribeRos2Message)
 from .aimrt_py_log import *
-from .aimrt_py_pb_chn import Publish, RegisterPublishType, Subscribe
 from .aimrt_python_runtime import *
diff --git a/...runtime/python_runtime/aimrt_py_pb_chn.py → src/runtime/python_runtime/aimrt_py_chn.py b/...runtime/python_runtime/aimrt_py_pb_chn.py → src/runtime/python_runtime/aimrt_py_chn.py
@@ -2,12 +2,15 @@
 # All rights reserved.
 
 import inspect
-from typing import Callable
+from typing import Callable, TypeVar
 
 import google.protobuf
 import google.protobuf.message
 
 from . import aimrt_python_runtime
+from .check_ros2_msg_type import check_is_valid_ros2_msg_type
+
+Ros2MsgType = TypeVar("Ros2MsgType")
 
 
 def _SerializeProtobufMessage(pb_msg: google.protobuf.message.Message, serialization_type: str) -> bytes:
@@ -32,21 +35,31 @@ def _DeserializeProtobufMessage(msg_buf: bytes,
     return msg
 
 
-def RegisterPublishType(publisher: aimrt_python_runtime.PublisherRef, protobuf_type: google.protobuf.message.Message):
+def RegisterPublishType(publisher: aimrt_python_runtime.PublisherRef,
+                        msg_type: google.protobuf.message.Message | Ros2MsgType):
     """Register a protobuf message type to a publisher.
 
     Args:
         publisher (aimrt_python_runtime.PublisherRef): channel publisher
-        protobuf_type (google.protobuf.message.Message): protobuf message type
+        msg_type (google.protobuf.message.Message | Ros2MsgType): protobuf message type or ROS2 message type
 
     Returns:
         bool: True if success, False otherwise
     """
-    aimrt_ts = aimrt_python_runtime.TypeSupport()
-    aimrt_ts.SetTypeName("pb:" + protobuf_type.DESCRIPTOR.full_name)
-    aimrt_ts.SetSerializationTypesSupportedList(["pb", "json"])
-
-    return publisher.RegisterPublishType(aimrt_ts)
+    if isinstance(msg_type, google.protobuf.message.Message):
+        aimrt_ts = aimrt_python_runtime.TypeSupport()
+        aimrt_ts.SetTypeName("pb:" + msg_type.DESCRIPTOR.full_name)
+        aimrt_ts.SetSerializationTypesSupportedList(["pb", "json"])
+        return publisher.RegisterPublishType(aimrt_ts)
+    elif check_is_valid_ros2_msg_type(msg_type):
+        # TODO: add ros2 type support
+        print(f"Register publish type for ROS2 message type: {msg_type.__name__}")
+        py_ros2_ts = aimrt_python_runtime.PyRos2TypeSupport(msg_type)
+        py_ros2_ts.SetTypeName("ros2:" + msg_type.__name__)
+        py_ros2_ts.SetSerializationTypesSupportedList(["ros2"])
+        return publisher.RegisterPublishType(py_ros2_ts)
+    else:
+        raise TypeError(f"Invalid message type: {type(msg_type)}")
 
 
 def Publish(publisher: aimrt_python_runtime.PublisherRef, second, third=None):
@@ -102,6 +115,13 @@ def Publish(publisher: aimrt_python_runtime.PublisherRef, second, third=None):
             f"only 'aimrt_python_runtime.Context' or 'aimrt_python_runtime.ContextRef' or 'str' is supported")
 
 
+def PublishRos2Message(publisher: aimrt_python_runtime.PublisherRef, msg: Ros2MsgType):
+    ctx = aimrt_python_runtime.Context()
+    ctx.SetSerializationType("ros2")
+    ctx_ref = aimrt_python_runtime.ContextRef(ctx)
+    publisher.PublishRos2MessageWithCtx("ros2:" + msg.__class__.__name__, ctx_ref, msg)
+
+
 def Subscribe(subscriber: aimrt_python_runtime.SubscriberRef,
               protobuf_type: google.protobuf.message.Message,
               callback: Callable):
@@ -140,3 +160,16 @@ def handle_callback(ctx_ref: aimrt_python_runtime.ContextRef, msg_buf: bytes):
             print(f"AimRT channel handle get exception, {e}")
 
     subscriber.SubscribeWithCtx(aimrt_ts, handle_callback)
+
+
+def SubscribeRos2Message(subscriber: aimrt_python_runtime.SubscriberRef,
+                         ros2_msg_type: Ros2MsgType,
+                         callback: Callable):
+    if not check_is_valid_ros2_msg_type(ros2_msg_type):
+        raise TypeError(f"Invalid ROS2 message type: {ros2_msg_type}")
+
+    py_ros2_ts = aimrt_python_runtime.PyRos2TypeSupport(ros2_msg_type)
+    py_ros2_ts.SetTypeName("ros2:" + ros2_msg_type.__name__)
+    py_ros2_ts.SetSerializationTypesSupportedList(["ros2"])
+
+    subscriber.SubscribeRos2MessageWithCtx(py_ros2_ts, ros2_msg_type, callback)
diff --git a/.../python_runtime/check_ro2_type_support.py → ...ime/python_runtime/check_ros2_msg_type.py b/.../python_runtime/check_ro2_type_support.py → ...ime/python_runtime/check_ros2_msg_type.py
@@ -1,19 +1,22 @@
 # Copyright (c) 2024 The AimRT Authors.
 # AimRT is licensed under Mulan PSL v2.
 
+import sys
+
+
 def check_for_ros2_type_support(msg_or_srv_type):
     try:
         ts = msg_or_srv_type.__class__._TYPE_SUPPORT
-    except AttributeError as e:
-        e.args = (
-            e.args[0] +
-            ' This might be a ROS 1 message type but it should be a ROS 2 message type.'
-            ' Make sure to source your ROS 2 workspace after your ROS 1 workspace.',
-            *e.args[1:])
+    except AttributeError:
+        print(
+            "This might be a ROS 1 message type but it should be a ROS 2 message type. "
+            "Make sure to source your ROS 2 workspace after your ROS 1 workspace.",
+            file=sys.stderr,
+        )
         return False
     if ts is None:
         msg_or_srv_type.__class__.__import_type_support__()
-    return msg_or_srv_type.__class__._TYPE_SUPPORT is None
+    return msg_or_srv_type.__class__._TYPE_SUPPORT is not None
 
 
 def check_is_valid_ros2_msg_type(msg_type):

diff --git a/src/runtime/python_runtime/export_channel.h b/src/runtime/python_runtime/export_channel.h
@@ -7,6 +7,8 @@
 
 #include "aimrt_module_cpp_interface/channel/channel_context.h"
 #include "aimrt_module_cpp_interface/channel/channel_handle.h"
+
+#include "python_runtime/export_ros2_type_support.h"
 #include "python_runtime/export_type_support.h"
 
 #include "pybind11/pybind11.h"
@@ -60,6 +62,15 @@ inline bool PyRegisterPublishType(
   return publisher_ref.RegisterPublishType(msg_type_support->NativeHandle());
 }
 
+inline bool PyRegisterPublishType(
+    aimrt::channel::PublisherRef& publisher_ref,
+    const std::shared_ptr<const PyRos2TypeSupport>& py_ros2_type_support) {
+  static std::vector<std::shared_ptr<const PyRos2TypeSupport>> py_ros2_ts_vec;
+  py_ros2_ts_vec.emplace_back(py_ros2_type_support);
+
+  return publisher_ref.RegisterPublishType(py_ros2_type_support->NativeHandle());
+}
+
 inline void PyPublishWithSerializationType(
     aimrt::channel::PublisherRef& publisher_ref,
     std::string_view msg_type,
@@ -78,15 +89,89 @@ inline void PyPublishWithCtx(
   publisher_ref.Publish(msg_type, ctx_ref, static_cast<const void*>(&msg_buf));
 }
 
+inline std::unique_ptr<void, destroy_ros_message_function*>
+create_from_py(pybind11::object pymessage) {
+  typedef void* create_ros_message_function(void);
+
+  py::object pymetaclass = pymessage.attr("__class__");
+
+  py::object value = pymetaclass.attr("_CREATE_ROS_MESSAGE");
+  auto capsule_ptr = static_cast<void*>(value.cast<py::capsule>());
+  auto create_ros_message =
+      reinterpret_cast<create_ros_message_function*>(capsule_ptr);
+  if (!create_ros_message) {
+    throw py::error_already_set();
+  }
+
+  value = pymetaclass.attr("_DESTROY_ROS_MESSAGE");
+  capsule_ptr = static_cast<void*>(value.cast<py::capsule>());
+  auto destroy_ros_message =
+      reinterpret_cast<destroy_ros_message_function*>(capsule_ptr);
+  if (!destroy_ros_message) {
+    throw py::error_already_set();
+  }
+
+  void* message = create_ros_message();
+  if (!message) {
+    throw std::bad_alloc();
+  }
+  return std::unique_ptr<
+      void, destroy_ros_message_function*>(message, destroy_ros_message);
+}
+
+inline std::unique_ptr<void, destroy_ros_message_function*>
+convert_from_py(py::object pymessage) {
+  typedef bool convert_from_py_signature(PyObject*, void*);
+
+  std::unique_ptr<void, destroy_ros_message_function*> message =
+      create_from_py(pymessage);
+
+  py::object pymetaclass = pymessage.attr("__class__");
+
+  auto capsule_ptr = static_cast<void*>(
+      pymetaclass.attr("_CONVERT_FROM_PY").cast<py::capsule>());
+  auto convert =
+      reinterpret_cast<convert_from_py_signature*>(capsule_ptr);
+  if (!convert) {
+    throw py::error_already_set();
+  }
+
+  if (!convert(pymessage.ptr(), message.get())) {
+    throw py::error_already_set();
+  }
+
+  return message;
+}
+
+inline void PyPublishRos2MessageWithCtx(
+    aimrt::channel::PublisherRef& publisher_ref,
+    std::string_view msg_type,
+    const aimrt::channel::ContextRef& ctx_ref,
+    pybind11::object msg_obj) {
+  auto msg_ptr = convert_from_py(msg_obj);
+  if (!msg_ptr) {
+    throw py::error_already_set();
+  }
+
+  publisher_ref.Publish(msg_type, ctx_ref, static_cast<const void*>(msg_ptr.get()));
+}
+
 inline void ExportPublisherRef(pybind11::object m) {
   using aimrt::channel::PublisherRef;
 
+  using PyTsPtr = std::shared_ptr<const PyTypeSupport>;
+  using PyRos2TsPtr = std::shared_ptr<const PyRos2TypeSupport>;
+
   pybind11::class_<PublisherRef>(std::move(m), "PublisherRef")
       .def(pybind11::init<>())
       .def("__bool__", &PublisherRef::operator bool)
-      .def("RegisterPublishType", &PyRegisterPublishType)
+      .def("RegisterPublishType",
+           pybind11::overload_cast<PublisherRef&, const PyTsPtr&>(&PyRegisterPublishType))
+      .def("RegisterPublishType",
+           pybind11::overload_cast<PublisherRef&, const PyRos2TsPtr&>(&PyRegisterPublishType))
       .def("PublishWithSerializationType", &PyPublishWithSerializationType)
       .def("PublishWithCtx", &PyPublishWithCtx)
+      .def("PublishRos2MessageWithCtx", &PyPublishRos2MessageWithCtx)
       .def("GetTopic", &PublisherRef::GetTopic)
       .def("MergeSubscribeContextToPublishContext", &PublisherRef::MergeSubscribeContextToPublishContext);
 }
@@ -121,13 +206,59 @@ inline bool PySubscribeWithCtx(
       });
 }
 
+inline py::object convert_to_py(void* message, py::object pyclass) {
+  py::object pymetaclass = pyclass.attr("__class__");
+
+  auto capsule_ptr = static_cast<void*>(
+      pymetaclass.attr("_CONVERT_TO_PY").cast<py::capsule>());
+
+  typedef PyObject* convert_to_py_function(void*);
+  auto convert = reinterpret_cast<convert_to_py_function*>(capsule_ptr);
+  if (!convert) {
+    throw py::error_already_set();
+  }
+  return py::reinterpret_steal<py::object>(convert(message));
+}
+
+inline bool PySubscribeRos2MessageWithCtx(
+    aimrt::channel::SubscriberRef& subscriber_ref,
+    const std::shared_ptr<const PyRos2TypeSupport>& py_ros2_type_support,
+    pybind11::object pyclass,
+    std::function<void(aimrt::channel::ContextRef, pybind11::object)>&& callback) {
+  static std::vector<std::shared_ptr<const PyRos2TypeSupport>> py_ros2_ts_vec;
+  py_ros2_ts_vec.emplace_back(py_ros2_type_support);
+
+  return subscriber_ref.Subscribe(
+      py_ros2_type_support->NativeHandle(),
+      [callback{std::move(callback)}, pyclass{pyclass}](
+          const aimrt_channel_context_base_t* ctx_ptr,
+          const void* msg_ptr,
+          aimrt_function_base_t* release_callback_base) {
+        aimrt::channel::SubscriberReleaseCallback release_callback(release_callback_base);
+
+        pybind11::gil_scoped_acquire acquire;
+        auto msg_obj = convert_to_py(const_cast<void*>(msg_ptr), pyclass);
+        if (!msg_obj) {
+          throw py::error_already_set();
+        }
+
+        auto ctx_ref = aimrt::channel::ContextRef(ctx_ptr);
+
+        callback(ctx_ref, msg_obj);
+        pybind11::gil_scoped_release release;
+
+        release_callback();
+      });
+}
+
 inline void ExportSubscriberRef(pybind11::object m) {
   using aimrt::channel::SubscriberRef;
 
   pybind11::class_<SubscriberRef>(std::move(m), "SubscriberRef")
       .def(pybind11::init<>())
       .def("__bool__", &SubscriberRef::operator bool)
       .def("SubscribeWithCtx", &PySubscribeWithCtx)
+      .def("SubscribeRos2MessageWithCtx", &PySubscribeRos2MessageWithCtx)
       .def("GetTopic", &SubscriberRef::GetTopic);
 }