Skip to content

Commit

Permalink
Add python bindings for the reader and writer by reusing the current
Browse files Browse the repository at this point in the history
Reader and Writer
  • Loading branch information
jmcarcell committed Jan 23, 2024
1 parent cd4ce5a commit 3e1f59b
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 4 deletions.
5 changes: 4 additions & 1 deletion python/podio/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,10 @@ def __init__(self, data=None):
"""
# Explicitly check for None here, to not return empty Frames on nullptr data
if data is not None:
self._frame = podio.Frame(data)
if isinstance(data, cppyy.gbl.podio.Frame):
self._frame = data
else:
self._frame = podio.Frame(data)
else:
self._frame = podio.Frame()

Expand Down
18 changes: 16 additions & 2 deletions python/podio/frame_iterator.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
# pylint: disable-next=import-error # gbl is a dynamic module from cppyy
from cppyy.gbl import std
from podio.frame import Frame
import cppyy


class FrameCategoryIterator:
Expand All @@ -27,7 +28,12 @@ def __iter__(self):

def __next__(self):
"""Get the next available Frame or stop."""
frame_data = self._reader.readNextEntry(self._category)
try:
frame_data = self._reader.readNextFrame(self._category)
except AttributeError:
frame_data = self._reader.readNextEntry(self._category)
except std.runtime_error:
raise StopIteration
if frame_data:
return Frame(std.move(frame_data))

Expand All @@ -52,14 +58,22 @@ def __getitem__(self, entry):
raise IndexError

try:
frame_data = self._reader.readFrame(self._category, entry)
except AttributeError:
frame_data = self._reader.readEntry(self._category, entry)
except std.bad_function_call:
print('Error: Unable to read an entry of the input file. This can happen when the '
'ROOT model dictionaries are not in LD_LIBRARY_PATH. Make sure that LD_LIBRARY_PATH '
'points to the library folder of the installation of podio and also to the library '
'folder with your data model\n')
raise
except std.runtime_error:
raise IndexError('Unable to read frame, you might be trying to read beyond bounds or a '
'non-existent category')

if frame_data:
return Frame(std.move(frame_data))
if isinstance(frame_data, cppyy.gbl.podio.ROOTFrameData):
return Frame(std.move(frame_data))
return Frame(frame_data)

raise IndexError
30 changes: 29 additions & 1 deletion python/podio/root_io.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
#!/usr/bin/env python3
"""Python module for reading root files containing podio Frames"""

from ROOT import gSystem
from ROOT import gSystem, gInterpreter
gSystem.Load('libpodioRootIO') # noqa: E402
gInterpreter.LoadFile("podio/Reader.h") # noqa: E402
from ROOT import podio # noqa: E402 # pylint: disable=wrong-import-position

from podio.base_reader import BaseReaderMixin # pylint: disable=wrong-import-position
Expand All @@ -12,6 +13,22 @@
class Reader(BaseReaderMixin):
"""Reader class for reading podio root files."""

def __init__(self, filenames):
"""Create a reader that reads from the passed file(s).
Args:
filenames (str or list[str]): file(s) to open and read data from
"""
if isinstance(filenames, str):
filenames = (filenames,)

self._reader = podio.makeReader(filenames)

super().__init__()

class ROOTFrameReader(BaseReaderMixin):
"""Reader class for reading podio root files."""

def __init__(self, filenames):
"""Create a reader that reads from the passed file(s).
Expand Down Expand Up @@ -69,6 +86,17 @@ def __init__(self, filenames):


class Writer(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.makeWriter(filename)
super().__init__()

class ROOTFrameWriter(BaseWriterMixin):
"""Writer class for writing podio root files"""
def __init__(self, filename):
"""Create a writer for writing files
Expand Down

0 comments on commit 3e1f59b

Please sign in to comment.