Python: raw_data + easier read_messages interface

lib/version.bzl

@@ -1 +1 @@
-EMBAG_VERSION = "0.0.41"
+EMBAG_VERSION = "0.0.42"

python/embag.cc

@@ -19,13 +19,26 @@ PYBIND11_MODULE(libembag, m) {
         return std::make_shared<Embag::Bag>(std::make_shared<const std::string>(bytes));
       }))
       .def("topics", &Embag::Bag::topics)
-      .def("read_messages", [](std::shared_ptr<Embag::Bag> &bag, const std::vector<std::string>& topics) {
-        Embag::View view{};
-        view.addBag(bag);
-        view.getMessages(topics);
+      .def(
+        "read_messages",
+        [](std::shared_ptr<Embag::Bag> &bag, py::object topics) {
+          Embag::View view{};
+          view.addBag(bag);
+          if (topics.is_none()) {
+            view.getMessages();
+          } else if (py::isinstance<py::str>(topics)) {
+            view.getMessages(topics.cast<std::string>());
+          } else if (py::isinstance<py::list>(topics)) {
+            view.getMessages(topics.cast<std::vector<std::string>>());
+          } else {
+            throw std::runtime_error("topics must be None, a string, or a list!");
+          }
 
-        return py::make_iterator(IteratorCompat{view.begin()}, IteratorCompat{view.end()});
-      }, py::keep_alive<0, 1>() /* Essential: keep object alive while iterator exists */, py::arg("topics"))
+          return py::make_iterator(IteratorCompat{view.begin()}, IteratorCompat{view.end()});
+        },
+        py::keep_alive<0, 1>(), /* Essential: keep object alive while iterator exists */
+        py::arg("topics") = py::none()
+      )
       .def("getSchema", [](std::shared_ptr<Embag::Bag> &bag, const std::string &topic) {
         auto builder = SchemaBuilder{bag};
         return builder.generateSchema(topic);
@@ -77,6 +90,9 @@ PYBIND11_MODULE(libembag, m) {
       .def("data", [](std::shared_ptr<Embag::RosMessage> &m) {
         return m->data();
       })
+      .def_property_readonly("raw_data", [](std::shared_ptr<Embag::RosMessage> &m) {
+        return py::bytes(&m->raw_buffer->at(m->raw_buffer_offset), m->raw_data_len);
+      })
       .def(
         "dict",
         [](

test/BUILD

@@ -16,6 +16,7 @@ py_test(
     srcs = ["embag_test_python.py"],
     data = [
         "test.bag",
+        "test_bag_raw_messages.P",
         "//python:libembag.so",
     ],
     python_version = "PY3",
@@ -27,6 +28,7 @@ py_test(
     srcs = ["embag_test_python.py"],
     data = [
         "test.bag",
+        "test_bag_raw_messages.P",
         "//python:libembag.so",
     ],
     python_version = "PY2",

test/embag_test_python.py

@@ -1,6 +1,7 @@
 import python.libembag as embag
 from collections import OrderedDict
 import numpy as np
+import pickle
 import struct
 import sys
 import unittest
@@ -243,5 +244,15 @@ def testROSTimeDicting(self):
             assert as_ros_time.to_nsec() == as_int
             assert as_ros_time.to_sec() == as_float
 
+    def testMessageRawData(self):
+        """
+        Tests that the raw_data property of a RosMessage returns the same data as rospy
+        """
+        with open('test/test_bag_raw_messages.P', 'rb') as raw_messages_file:
+            raw_messages = pickle.load(raw_messages_file)
+
+        for index, msg in enumerate(self.view.getMessages()):
+            assert raw_messages[index] == msg.raw_data
+
 if __name__ == "__main__":
     unittest.main()