Add python bindings for RNTuple writer

Author: tmadlener

python/podio/root_io.py

@@ -77,3 +77,14 @@ def __init__(self, filename):
         filename (str): The name of the output file
     """
     self._writer = podio.ROOTFrameWriter(filename)
+
+
+class RNTupleWriter(BaseWriterMixin):
+  """Writer class for writing podio root files"""
+  def __init__(self, filename):
+    """Create a writer for writing files
+
+    Args:
+        filename (str): The name of the output file
+    """
+    self._writer = podio.ROOTNTupleWriter(filename)

tests/root_io/CMakeLists.txt

@@ -19,6 +19,7 @@ if(ENABLE_RNTUPLE)
       ${root_dependent_tests}
       write_rntuple.cpp
       read_rntuple.cpp
+      read_python_frame_rntuple.cpp
      )
 endif()
 set(root_libs TestDataModelDict ExtensionDataModelDict podio::podioRootIO)
@@ -80,6 +81,13 @@ endforeach()
 
 #--- Write via python and the ROOT backend and see if we can read it back in in
 #--- c++
-add_test(NAME write_python_frame_root COMMAND python3 ${PROJECT_SOURCE_DIR}/tests/write_frame.py example_frame_with_py.root)
+add_test(NAME write_python_frame_root COMMAND python3 ${PROJECT_SOURCE_DIR}/tests/write_frame.py example_frame_with_py.root root_io.Writer)
 PODIO_SET_TEST_ENV(write_python_frame_root)
 set_property(TEST read_python_frame_root PROPERTY DEPENDS write_python_frame_root)
+
+if (ENABLE_RNTUPLE)
+  add_test(NAME write_python_frame_rntuple COMMAND python3 ${PROJECT_SOURCE_DIR}/tests/write_frame.py example_frame_with_py_rntuple.root root_io.RNTupleWriter)
+  PODIO_SET_TEST_ENV(write_python_frame_rntuple)
+
+  set_property(TEST read_python_frame_rntuple PROPERTY DEPENDS write_python_frame_rntuple)
+endif()

tests/root_io/read_python_frame_rntuple.cpp

@@ -0,0 +1,7 @@
+#include "read_python_frame.h"
+
+#include "podio/ROOTNTupleReader.h"
+
+int main() {
+  return read_frame<podio::ROOTNTupleReader>("example_frame_with_py_rntuple.root");
+}

tests/sio_io/CMakeLists.txt

@@ -40,6 +40,6 @@ set_property(TEST check_benchmark_outputs_sio PROPERTY DEPENDS read_timed_sio wr
 
 #--- Write via python and the SIO backend and see if we can read it back in in
 #--- c++
-add_test(NAME write_python_frame_sio COMMAND python3 ${PROJECT_SOURCE_DIR}/tests/write_frame.py example_frame_with_py.sio)
+add_test(NAME write_python_frame_sio COMMAND python3 ${PROJECT_SOURCE_DIR}/tests/write_frame.py example_frame_with_py.sio sio_io.Writer)
 PODIO_SET_TEST_ENV(write_python_frame_sio)
 set_property(TEST read_python_frame_sio PROPERTY DEPENDS write_python_frame_sio)

tests/write_frame.py

@@ -54,13 +54,14 @@ def create_frame():
   return frame
 
 
-def write_file(io_backend, filename):
+def write_file(writer_type, filename):
   """Write a file using the given Writer type and put one Frame into it under
   the events category
   """
+  io_backend, writer_name = writer_type.split(".")
   io_module = importlib.import_module(f"podio.{io_backend}")
 
-  writer = io_module.Writer(filename)
+  writer = getattr(io_module, writer_name)(filename)
   event = create_frame()
   writer.write_frame(event, "events")
 
@@ -70,9 +71,10 @@ def write_file(io_backend, filename):
 
   parser = argparse.ArgumentParser()
   parser.add_argument("outputfile", help="Output file name")
+  parser.add_argument("writer", help="The writer type to use")
 
   args = parser.parse_args()
 
   io_format = args.outputfile.split(".")[-1]
 
-  write_file(f"{io_format}_io", args.outputfile)
+  write_file(args.writer, args.outputfile)