Skip to content

Commit

Permalink
Add python bindings for RNTuple writer
Browse files Browse the repository at this point in the history
  • Loading branch information
tmadlener committed Nov 24, 2023
1 parent 5b5f911 commit fa559fd
Show file tree
Hide file tree
Showing 5 changed files with 33 additions and 5 deletions.
11 changes: 11 additions & 0 deletions python/podio/root_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
10 changes: 9 additions & 1 deletion tests/root_io/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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()
7 changes: 7 additions & 0 deletions tests/root_io/read_python_frame_rntuple.cpp
Original file line number Diff line number Diff line change
@@ -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");
}
2 changes: 1 addition & 1 deletion tests/sio_io/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)
8 changes: 5 additions & 3 deletions tests/write_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand All @@ -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)

0 comments on commit fa559fd

Please sign in to comment.