expose py Reader metadata, improve rosbag2_py.BagMetadata usability (…

…#1082) * expose metadata for reader, fix metadata fields to yield rclpy.Time and rclpy.Duration values instead of datetime * Provide default constructor arguments for rosbag2_py.BagMetadata * Add tests for rosbag2_py Storage struct constructors, fixing found issues * Expose all StorageFilter fields in py Signed-off-by: Emerson Knapp <[email protected]>
ros2 · Sep 29, 2022 · c58d73c · c58d73c
1 parent 6de3611
commit c58d73c
Show file tree

Hide file tree

Showing 12 changed files with 313 additions and 140 deletions.
diff --git a/rosbag2_py/CMakeLists.txt b/rosbag2_py/CMakeLists.txt
@@ -118,12 +118,13 @@ if(BUILD_TESTING)
     set(set_env_vars "${set_env_vars} ROSBAG2_PY_TEST_WITH_RTLD_GLOBAL=True")
   endif()
 
-  # message(FATAL_ERROR "HELP I AM THE PATH ${CMAKE_CURRENT_SOURCE_DIR}")
-  ament_add_pytest_test(test_sequential_reader_py "test/test_sequential_reader.py"
+  ament_add_pytest_test(test_sequential_reader_py
+    "test/test_sequential_reader.py"
     APPEND_ENV "${append_env_vars}"
     ENV "${set_env_vars}"
   )
-  ament_add_pytest_test(test_sequential_reader_multiple_files_py "test/test_sequential_reader_multiple_files.py"
+  ament_add_pytest_test(test_sequential_reader_multiple_files_py
+    "test/test_sequential_reader_multiple_files.py"
     APPEND_ENV "${append_env_vars}"
     ENV "${set_env_vars}"
   )
@@ -147,6 +148,12 @@ if(BUILD_TESTING)
     APPEND_ENV "${append_env_vars}"
     ENV "${set_env_vars}"
   )
+
+  ament_add_pytest_test(test_storage_py
+    "test/test_storage.py"
+    APPEND_ENV "${append_env_vars}"
+    ENV "${set_env_vars}"
+  )
 endif()
 
 ament_package()
diff --git a/rosbag2_py/rosbag2_py/__init__.py b/rosbag2_py/rosbag2_py/__init__.py
@@ -25,6 +25,7 @@
     )
     from rosbag2_py._storage import (
         ConverterOptions,
+        FileInformation,
         StorageFilter,
         StorageOptions,
         TopicMetadata,
@@ -55,6 +56,7 @@
 __all__ = [
     'bag_rewrite',
     'ConverterOptions',
+    'FileInformation',
     'get_registered_readers',
     'get_registered_writers',
     'get_registered_compressors',

diff --git a/rosbag2_py/src/rosbag2_py/_reader.cpp b/rosbag2_py/src/rosbag2_py/_reader.cpp
@@ -45,7 +45,7 @@ class Reader
 
   void open(
     rosbag2_storage::StorageOptions & storage_options,
-    rosbag2_cpp::ConverterOptions & converter_options)
+    rosbag2_cpp::ConverterOptions & converter_options = rosbag2_cpp::ConverterOptions())
   {
     reader_->open(storage_options, converter_options);
   }
@@ -73,6 +73,11 @@ class Reader
     return reader_->get_all_topics_and_types();
   }
 
+  const rosbag2_storage::BagMetadata & get_metadata() const
+  {
+    return reader_->get_metadata();
+  }
+
   void set_filter(const rosbag2_storage::StorageFilter & storage_filter)
   {
     return reader_->set_filter(storage_filter);
@@ -108,44 +113,33 @@ std::unordered_set<std::string> get_registered_readers()
 
 }  // namespace rosbag2_py
 
+using PyReader = rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>;
+using PyCompressionReader = rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>;
+
 PYBIND11_MODULE(_reader, m) {
   m.doc() = "Python wrapper of the rosbag2_cpp reader API";
 
-  pybind11::class_<rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>>(
-    m, "SequentialReader")
+  pybind11::class_<PyReader>(m, "SequentialReader")
   .def(pybind11::init())
-  .def("open", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::open)
-  .def("read_next", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::read_next)
-  .def("has_next", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::has_next)
-  .def(
-    "get_all_topics_and_types",
-    &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::get_all_topics_and_types)
-  .def("set_filter", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::set_filter)
-  .def("reset_filter", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::reset_filter)
-  .def("seek", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::seek);
-
-  pybind11::class_<rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>>(
-    m, "SequentialCompressionReader")
+  .def("open", &PyReader::open)
+  .def("read_next", &PyReader::read_next)
+  .def("has_next", &PyReader::has_next)
+  .def("get_metadata", &PyReader::get_metadata)
+  .def("get_all_topics_and_types", &PyReader::get_all_topics_and_types)
+  .def("set_filter", &PyReader::set_filter)
+  .def("reset_filter", &PyReader::reset_filter)
+  .def("seek", &PyReader::seek);
+
+  pybind11::class_<PyCompressionReader>(m, "SequentialCompressionReader")
   .def(pybind11::init())
-  .def("open", &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::open)
-  .def(
-    "read_next", &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::read_next)
-  .def("has_next", &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::has_next)
-  .def(
-    "get_all_topics_and_types",
-    &rosbag2_py::Reader<
-      rosbag2_compression::SequentialCompressionReader
-    >::get_all_topics_and_types)
-  .def(
-    "set_filter",
-    &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::set_filter)
-  .def(
-    "reset_filter",
-    &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::reset_filter)
-  .def(
-    "seek",
-    &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::seek);
-
+  .def("open", &PyCompressionReader::open)
+  .def("read_next", &PyCompressionReader::read_next)
+  .def("has_next", &PyCompressionReader::has_next)
+  .def("get_metadata", &PyCompressionReader::get_metadata)
+  .def("get_all_topics_and_types", &PyCompressionReader::get_all_topics_and_types)
+  .def("set_filter", &PyCompressionReader::set_filter)
+  .def("reset_filter", &PyCompressionReader::reset_filter)
+  .def("seek", &PyCompressionReader::seek);
   m.def(
     "get_registered_readers",
     &rosbag2_py::get_registered_readers,

diff --git a/rosbag2_py/src/rosbag2_py/_storage.cpp b/rosbag2_py/src/rosbag2_py/_storage.cpp
@@ -25,14 +25,50 @@
 
 #include "./pybind11.hpp"
 
+namespace
+{
+
+using pybind11::literals::operator""_a;
+
+pybind11::object to_rclpy_duration(std::chrono::nanoseconds duration)
+{
+  pybind11::object Duration = pybind11::module::import("rclpy.duration").attr("Duration");
+  return Duration("nanoseconds"_a = duration.count());
+}
+
+std::chrono::nanoseconds from_rclpy_duration(const pybind11::object & duration)
+{
+  pybind11::int_ nanos = duration.attr("nanoseconds");
+  return std::chrono::nanoseconds(nanos);
+}
+
+template<typename T>
+pybind11::object to_rclpy_time(T time)
+{
+  pybind11::object Time = pybind11::module::import("rclpy.time").attr("Time");
+  return Time(
+    "nanoseconds"_a = std::chrono::duration_cast<std::chrono::nanoseconds>(
+      time.time_since_epoch()).count());
+}
+
+std::chrono::time_point<std::chrono::high_resolution_clock> from_rclpy_time(
+  const pybind11::object & time)
+{
+  pybind11::int_ nanos = time.attr("nanoseconds");
+  return std::chrono::time_point<std::chrono::high_resolution_clock>(
+    std::chrono::nanoseconds(nanos));
+}
+
+}  // namespace
+
 PYBIND11_MODULE(_storage, m) {
   m.doc() = "Python wrapper of the rosbag2 utilities API";
 
   pybind11::class_<rosbag2_cpp::ConverterOptions>(m, "ConverterOptions")
   .def(
     pybind11::init<std::string, std::string>(),
-    pybind11::arg("input_serialization_format"),
-    pybind11::arg("output_serialization_format"))
+    pybind11::arg("input_serialization_format") = "",
+    pybind11::arg("output_serialization_format") = "")
   .def_readwrite(
     "input_serialization_format",
     &rosbag2_cpp::ConverterOptions::input_serialization_format)
@@ -81,9 +117,15 @@ PYBIND11_MODULE(_storage, m) {
 
   pybind11::class_<rosbag2_storage::StorageFilter>(m, "StorageFilter")
   .def(
-    pybind11::init<std::vector<std::string>>(),
-    pybind11::arg("topics"))
-  .def_readwrite("topics", &rosbag2_storage::StorageFilter::topics);
+    pybind11::init<std::vector<std::string>, std::string, std::string>(),
+    pybind11::arg("topics") = std::vector<std::string>(),
+    pybind11::arg("topics_regex") = "",
+    pybind11::arg("topics_regex_to_exclude") = "")
+  .def_readwrite("topics", &rosbag2_storage::StorageFilter::topics)
+  .def_readwrite("topics_regex", &rosbag2_storage::StorageFilter::topics_regex)
+  .def_readwrite(
+    "topics_regex_to_exclude",
+    &rosbag2_storage::StorageFilter::topics_regex_to_exclude);
 
   pybind11::class_<rosbag2_storage::TopicMetadata>(m, "TopicMetadata")
   .def(
@@ -112,57 +154,118 @@ PYBIND11_MODULE(_storage, m) {
 
   pybind11::class_<rosbag2_storage::FileInformation>(m, "FileInformation")
   .def(
-    pybind11::init<std::string,
-    std::chrono::time_point<std::chrono::high_resolution_clock>,
-    std::chrono::nanoseconds,
-    size_t>(),
+    pybind11::init(
+      [](
+        std::string path,
+        pybind11::object starting_time,
+        pybind11::object duration,
+        size_t message_count)
+      {
+        return rosbag2_storage::FileInformation{
+          path,
+          from_rclpy_time(starting_time),
+          from_rclpy_duration(duration),
+          message_count
+        };
+      }),
     pybind11::arg("path"),
     pybind11::arg("starting_time"),
     pybind11::arg("duration"),
-    pybind11::arg("message_count"))
+    pybind11::arg("message_count")
+  )
   .def_readwrite("path", &rosbag2_storage::FileInformation::path)
-  .def_readwrite("starting_time", &rosbag2_storage::FileInformation::starting_time)
+  .def_property(
+    "starting_time",
+    [](const rosbag2_storage::FileInformation & self) {
+      return to_rclpy_time(self.starting_time);
+    },
+    [](rosbag2_storage::FileInformation & self, const pybind11::object & value) {
+      self.starting_time = from_rclpy_time(value);
+    })
+  .def_property(
+    "duration",
+    [](const rosbag2_storage::FileInformation & self) {
+      return to_rclpy_duration(self.duration);
+    },
+    [](rosbag2_storage::FileInformation & self, const pybind11::object & value) {
+      self.duration = from_rclpy_duration(value);
+    })
   .def_readwrite("duration", &rosbag2_storage::FileInformation::duration)
   .def_readwrite("message_count", &rosbag2_storage::FileInformation::message_count);
 
   pybind11::class_<rosbag2_storage::BagMetadata>(m, "BagMetadata")
   .def(
-    pybind11::init<
-      int,
-      uint64_t,
-      std::string,
-      std::vector<std::string>,
-      std::vector<rosbag2_storage::FileInformation>,
-      std::chrono::nanoseconds,
-      std::chrono::time_point<std::chrono::high_resolution_clock>,
-      uint64_t,
-      std::vector<rosbag2_storage::TopicInformation>,
-      std::string,
-      std::string>(),
-    pybind11::arg("version"),
-    pybind11::arg("bag_size"),
-    pybind11::arg("storage_identifier"),
-    pybind11::arg("relative_file_paths"),
-    pybind11::arg("files"),
-    pybind11::arg("duration"),
-    pybind11::arg("starting_time"),
-    pybind11::arg("message_count"),
-    pybind11::arg("topics_with_message_count"),
-    pybind11::arg("compression_format"),
-    pybind11::arg("compression_mode"))
+    pybind11::init(
+      [](
+        int version,
+        uint64_t bag_size,
+        std::string storage_identifier,
+        std::vector<std::string> relative_file_paths,
+        std::vector<rosbag2_storage::FileInformation> files,
+        pybind11::object duration,
+        pybind11::object starting_time,
+        uint64_t message_count,
+        std::vector<rosbag2_storage::TopicInformation> topics_with_message_count,
+        std::string compression_format,
+        std::string compression_mode,
+        std::unordered_map<std::string, std::string> custom_data)
+      {
+        return rosbag2_storage::BagMetadata{
+          version,
+          bag_size,
+          storage_identifier,
+          relative_file_paths,
+          files,
+          from_rclpy_duration(duration),
+          from_rclpy_time(starting_time),
+          message_count,
+          topics_with_message_count,
+          compression_format,
+          compression_mode,
+          custom_data
+        };
+      }),
+    pybind11::arg("version") = 6,
+    pybind11::arg("bag_size") = 0,
+    pybind11::arg("storage_identifier") = "",
+    pybind11::arg("relative_file_paths") = std::vector<std::string>(),
+    pybind11::arg("files") = std::vector<rosbag2_storage::FileInformation>(),
+    pybind11::arg("duration") = to_rclpy_duration(std::chrono::nanoseconds{0}),
+    pybind11::arg("starting_time") = to_rclpy_time(
+      std::chrono::time_point<std::chrono::high_resolution_clock>(std::chrono::nanoseconds{0})),
+    pybind11::arg("message_count") = 0,
+    pybind11::arg("topics_with_message_count") = std::vector<rosbag2_storage::TopicInformation>(),
+    pybind11::arg("compression_format") = "",
+    pybind11::arg("compression_mode") = "",
+    pybind11::arg("custom_data") = std::unordered_map<std::string, std::string>())
   .def_readwrite("version", &rosbag2_storage::BagMetadata::version)
   .def_readwrite("bag_size", &rosbag2_storage::BagMetadata::bag_size)
   .def_readwrite("storage_identifier", &rosbag2_storage::BagMetadata::storage_identifier)
   .def_readwrite("relative_file_paths", &rosbag2_storage::BagMetadata::relative_file_paths)
   .def_readwrite("files", &rosbag2_storage::BagMetadata::files)
-  .def_readwrite("duration", &rosbag2_storage::BagMetadata::duration)
-  .def_readwrite("starting_time", &rosbag2_storage::BagMetadata::starting_time)
+  .def_property(
+    "duration",
+    [](const rosbag2_storage::BagMetadata & self) {
+      return to_rclpy_duration(self.duration);
+    },
+    [](rosbag2_storage::BagMetadata & self, const pybind11::object & value) {
+      self.duration = from_rclpy_duration(value);
+    })
+  .def_property(
+    "starting_time",
+    [](const rosbag2_storage::BagMetadata & self) {
+      return to_rclpy_time(self.starting_time);
+    },
+    [](rosbag2_storage::BagMetadata & self, const pybind11::object & value) {
+      self.starting_time = from_rclpy_time(value);
+    })
   .def_readwrite("message_count", &rosbag2_storage::BagMetadata::message_count)
   .def_readwrite(
     "topics_with_message_count",
     &rosbag2_storage::BagMetadata::topics_with_message_count)
   .def_readwrite("compression_format", &rosbag2_storage::BagMetadata::compression_format)
   .def_readwrite("compression_mode", &rosbag2_storage::BagMetadata::compression_mode)
+  .def_readwrite("custom_data", &rosbag2_storage::BagMetadata::custom_data)
   .def(
     "__repr__", [](const rosbag2_storage::BagMetadata & metadata) {
       return format_bag_meta_data(metadata);

diff --git a/rosbag2_py/test/common.py b/rosbag2_py/test/common.py
@@ -12,12 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import sys
 import time
 from typing import Callable
 
-from rclpy.clock import Clock, ClockType
-from rclpy.duration import Duration
-import rosbag2_py
+if os.environ.get('ROSBAG2_PY_TEST_WITH_RTLD_GLOBAL', None) is not None:
+    # This is needed on Linux when compiling with clang/libc++.
+    # TL;DR This makes class_loader work when using a python extension compiled with libc++.
+    #
+    # For the fun RTTI ABI details, see https://whatofhow.wordpress.com/2015/03/17/odr-rtti-dso/.
+    sys.setdlopenflags(os.RTLD_GLOBAL | os.RTLD_LAZY)
+
+from rclpy.clock import Clock, ClockType  # noqa
+from rclpy.duration import Duration  # noqa
+import rosbag2_py  # noqa
 
 
 def get_rosbag_options(path, serialization_format='cdr'):