diff --git a/rosbag2_py/CMakeLists.txt b/rosbag2_py/CMakeLists.txt index 91831730ac..40cba72e89 100644 --- a/rosbag2_py/CMakeLists.txt +++ b/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}" diff --git a/rosbag2_py/rosbag2_py/__init__.py b/rosbag2_py/rosbag2_py/__init__.py index 72339859da..64c0374ce3 100644 --- a/rosbag2_py/rosbag2_py/__init__.py +++ b/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', diff --git a/rosbag2_py/src/rosbag2_py/_compression_options.cpp b/rosbag2_py/src/rosbag2_py/_compression_options.cpp new file mode 100644 index 0000000000..bf471ed208 --- /dev/null +++ b/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 +#include + +#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_(m, "CompressionMode") + .value("NONE", CompressionMode::NONE) + .value("FILE", CompressionMode::FILE) + .value("MESSAGE", CompressionMode::MESSAGE) + .export_values(); + + pybind11::class_(m, "CompressionOptions") + .def(pybind11::init(), + 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"); + +} diff --git a/rosbag2_py/src/rosbag2_py/_writer.cpp b/rosbag2_py/src/rosbag2_py/_writer.cpp index eeb29a25df..1841e1eddd 100644 --- a/rosbag2_py/src/rosbag2_py/_writer.cpp +++ b/rosbag2_py/src/rosbag2_py/_writer.cpp @@ -38,8 +38,9 @@ template class Writer : public rosbag2_cpp::Writer { public: - Writer() - : rosbag2_cpp::Writer(std::make_unique()) + template + Writer(Args &&... args) + : rosbag2_cpp::Writer(std::make_unique(std::forward(args)...)) {} /// Write a serialized message to a bag file @@ -104,7 +105,7 @@ PYBIND11_MODULE(_writer, m) { ; pybind11::class_(m, "SequentialCompressionWriter") - .def(pybind11::init()) + .def(pybind11::init()) .def( "open", pybind11::overload_cast< diff --git a/rosbag2_py/test/test_compression.py b/rosbag2_py/test/test_compression.py new file mode 100644 index 0000000000..7dfe2e7185 --- /dev/null +++ b/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