Add support for compression to python API

rosbag2_py/CMakeLists.txt

@@ -40,6 +40,13 @@ find_package(pybind11 REQUIRED)
 
 ament_python_install_package(${PROJECT_NAME})
 
+pybind11_add_module(_compression_options SHARED
+  src/rosbag2_py/_compression_options.cpp
+)
+ament_target_dependencies(_compression_options PUBLIC
+  "rosbag2_compression"
+)
+
 pybind11_add_module(_reader SHARED
   src/rosbag2_py/_reader.cpp
 )
@@ -96,6 +103,7 @@ ament_target_dependencies(_reindexer PUBLIC
 # Install cython modules as sub-modules of the project
 install(
   TARGETS
+    _compression_options
     _reader
     _storage
     _writer
@@ -148,7 +156,11 @@ if(BUILD_TESTING)
     APPEND_ENV "${append_env_vars}"
     ENV "${set_env_vars}"
   )
-
+  ament_add_pytest_test(test_compression_py
+    "test/test_compression.py"
+    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}"

rosbag2_py/rosbag2_py/__init__.py

@@ -18,6 +18,12 @@
 # to the search path.
 # See https://docs.python.org/3/whatsnew/3.8.html#bpo-36085-whatsnew
 with add_dll_directories_from_env('PATH'):
+    from rosbag2_py._compression_options import (
+        CompressionMode,
+        CompressionOptions,
+        compression_mode_from_string,
+        compression_mode_to_string
+    )
     from rosbag2_py._reader import (
         SequentialCompressionReader,
         SequentialReader,
@@ -59,6 +65,10 @@
 
 __all__ = [
     'bag_rewrite',
+    'CompressionMode',
+    'CompressionOptions',
+    'compression_mode_from_string',
+    'compression_mode_to_string',
     'ConverterOptions',
     'FileInformation',
     'get_default_storage_id',

rosbag2_py/src/rosbag2_py/_compression_options.cpp

@@ -0,0 +1,55 @@
+// Copyright 2023 Open Source Robotics Foundation, Inc.
+//
+// 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.
+
+#include <memory>
+#include <string>
+
+#include "rosbag2_compression/compression_options.hpp"
+
+#include "./pybind11.hpp"
+
+using CompressionMode = rosbag2_compression::CompressionMode;
+using CompressionOptions = rosbag2_compression::CompressionOptions;
+
+PYBIND11_MODULE(_compression_options, m) {
+  m.doc() = "Python wrapper of the rosbag2_compression API";
+
+  pybind11::enum_<CompressionMode>(m, "CompressionMode")
+  .value("NONE", CompressionMode::NONE)
+  .value("FILE", CompressionMode::FILE)
+  .value("MESSAGE", CompressionMode::MESSAGE)
+  .export_values();
+
+  pybind11::class_<CompressionOptions>(m, "CompressionOptions")
+  .def(pybind11::init<std::string &, CompressionMode, uint64_t &, uint64_t &>(),
+      pybind11::arg("compression_format") = CompressionOptions{}.compression_format,
+      pybind11::arg("compression_mode") = CompressionOptions{}.compression_mode,
+      pybind11::arg("compression_queue_size") = CompressionOptions{}.compression_queue_size,
+      pybind11::arg("compression_threads") = CompressionOptions{}.compression_threads)
+  .def_readwrite("compression_format", &CompressionOptions::compression_format)
+  .def_readwrite("compression_mode", &CompressionOptions::compression_mode)
+  .def_readwrite("compression_queue_size", &CompressionOptions::compression_queue_size)
+  .def_readwrite("compression_threads", &CompressionOptions::compression_threads);
+
+  m.def(
+    "compression_mode_from_string",
+    &rosbag2_compression::compression_mode_from_string,
+    "Converts a string into a rosbag2_compression::CompressionMode enum.");
+
+  m.def(
+    "compression_mode_to_string",
+    &rosbag2_compression::compression_mode_to_string,
+    "Converts a rosbag2_compression::CompressionMode enum into a string");
+
+}

rosbag2_py/src/rosbag2_py/_writer.cpp

@@ -38,8 +38,9 @@ template<typename T>
 class Writer : public rosbag2_cpp::Writer
 {
 public:
-  Writer()
-  : rosbag2_cpp::Writer(std::make_unique<T>())
+  template <typename... Args>
+  Writer(Args &&... args)
+  : rosbag2_cpp::Writer(std::make_unique<T>(std::forward<Args>(args)...))
   {}
 
   /// Write a serialized message to a bag file
@@ -104,7 +105,7 @@ PYBIND11_MODULE(_writer, m) {
   ;
 
   pybind11::class_<PyCompressionWriter>(m, "SequentialCompressionWriter")
-  .def(pybind11::init())
+  .def(pybind11::init<rosbag2_compression::CompressionOptions>())
   .def(
     "open",
     pybind11::overload_cast<

rosbag2_py/test/test_compression.py

@@ -0,0 +1,101 @@
+# Copyright 2022, Foxglove Technologies. All rights reserved.
+#
+# 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.
+
+import os
+import unittest
+
+from rosidl_runtime_py.utilities import get_message
+from rclpy.duration import Duration
+from rclpy.serialization import deserialize_message, serialize_message
+from rclpy.time import Time
+from rosbag2_py import (
+    CompressionMode,
+    CompressionOptions,
+    compression_mode_from_string,
+    compression_mode_to_string,
+    SequentialCompressionReader,
+    SequentialCompressionWriter,
+    TopicMetadata
+)
+
+from std_msgs.msg import String
+
+from common import get_rosbag_options
+
+
+def test_compression_mode_from_string():
+    """Checks that we can cast a string to a compression mode"""
+
+    assert CompressionMode.NONE == compression_mode_from_string("NONE")
+    assert CompressionMode.MESSAGE == compression_mode_from_string("MESSAGE")
+    assert CompressionMode.FILE == compression_mode_from_string("FILE")
+
+def test_compression_mode_to_string():
+    """Checks that we can cast a compression mode to a string"""
+
+    assert "NONE" == compression_mode_to_string(CompressionMode.NONE)
+    assert "MESSAGE" == compression_mode_to_string(CompressionMode.MESSAGE)
+    assert "FILE" == compression_mode_to_string(CompressionMode.FILE)
+
+def test_compression_options():
+    """Checks that we can construct a CompressionOptions class"""
+
+    compression_options = CompressionOptions(
+        compression_format = "zstd",
+        compression_mode = CompressionMode.MESSAGE,
+        compression_queue_size = 0,
+        compression_threads = 8)
+    assert compression_options is not None
+    assert compression_options.compression_format == "zstd"
+    assert compression_options.compression_mode == CompressionMode.MESSAGE
+    assert compression_options.compression_queue_size == 0
+    assert compression_options.compression_threads == 8
+
+
+def test_sequential_compression_writer(tmp_path):
+    """Checks that we can do a compressed write and read"""    
+
+    bag_path = os.path.join(tmp_path, 'tmp_sequential_compressed_write_test')
+
+    storage_id = 'mcap'
+
+    storage_options, converter_options = get_rosbag_options(
+        path = bag_path,
+        storage_id = storage_id)
+
+    compression_options = CompressionOptions(
+        compression_format = 'zstd',
+        compression_mode = CompressionMode.MESSAGE,
+        compression_queue_size = 0,
+        compression_threads = 16)
+
+    writer = SequentialCompressionWriter(compression_options)
+    writer.open(storage_options, converter_options)
+
+    topic_name = '/chatter'
+    topic_metadata = TopicMetadata(
+        name=topic_name,
+        type='std_msgs/msg/String',
+        serialization_format='cdr')
+    writer.create_topic(topic_metadata)
+
+    for i in range(10):
+        msg = String()
+        msg.data = f'Hello, world! {str(i)}'
+        time_stamp = i * 100
+
+        writer.write(topic_name, serialize_message(msg), time_stamp)
+
+    # close bag and create new storage instance
+    del writer