Add support for compression to python API

rosbag2_py/CMakeLists.txt

@@ -41,6 +41,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
 )
@@ -97,6 +104,7 @@ target_link_libraries(_reindexer PUBLIC
 # Install cython modules as sub-modules of the project
 install(
   TARGETS
+    _compression_options
     _reader
     _storage
     _writer
@@ -149,7 +157,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/package.xml

@@ -31,6 +31,7 @@
   <test_depend>python3-pytest</test_depend>
   <test_depend>rcl_interfaces</test_depend>
   <test_depend>rclpy</test_depend>
+  <test_depend>rosbag2_compression_zstd</test_depend>
   <test_depend>rosbag2_storage_default_plugins</test_depend>
   <test_depend>rosbag2_test_common</test_depend>
   <test_depend>rosidl_runtime_py</test_depend>

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,
@@ -60,6 +66,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>
+  explicit Writer(Args && ... args)
+  : rosbag2_cpp::Writer(std::make_unique<T>(std::forward<Args>(args)...))
   {}
 
   /// Write a serialized message to a bag file
@@ -107,7 +108,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,120 @@
+# 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
+
+from common import get_rosbag_options
+
+import pytest
+
+from rclpy.serialization import deserialize_message, serialize_message
+from rosbag2_py import (
+    compression_mode_from_string,
+    compression_mode_to_string,
+    CompressionMode,
+    CompressionOptions,
+    SequentialCompressionReader,
+    SequentialCompressionWriter,
+    TopicMetadata,
+)
+from rosbag2_test_common import TESTED_STORAGE_IDS
+from rosidl_runtime_py.utilities import get_message
+
+from std_msgs.msg import String
+
+
+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
+
+
+@pytest.mark.parametrize('storage_id', TESTED_STORAGE_IDS)
+def test_sequential_compression(tmp_path, storage_id):
+    """Checks that we can do a compressed write and read."""
+    bag_path = os.path.join(tmp_path, 'tmp_sequential_compressed_write_test')
+
+    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=1)
+
+    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
+
+    storage_options, converter_options = get_rosbag_options(bag_path, storage_id)
+
+    reader = SequentialCompressionReader()
+    reader.open(storage_options, converter_options)
+
+    topic_types = reader.get_all_topics_and_types()
+
+    # Create a map for quicker lookup
+    type_map = {topic_types[i].name: topic_types[i].type for i in range(len(topic_types))}
+
+    msg_counter = 0
+    while reader.has_next():
+        topic, data, t = reader.read_next()
+        msg_type = get_message(type_map[topic])
+        msg_deserialized = deserialize_message(data, msg_type)
+
+        assert isinstance(msg_deserialized, String)
+        assert msg_deserialized.data == f'Hello, world! {msg_counter}'
+        assert t == msg_counter * 100
+
+        msg_counter += 1