From 3e1f59bd7f3c0fa5947dc219f2e6593257f1be46 Mon Sep 17 00:00:00 2001
From: jmcarcell <jmcarcell@users.noreply.github.com>
Date: Tue, 23 Jan 2024 08:51:35 +0100
Subject: [PATCH] Add python bindings for the reader and writer by reusing the
 current

Reader and Writer
---
 python/podio/frame.py          |  5 ++++-
 python/podio/frame_iterator.py | 18 ++++++++++++++++--
 python/podio/root_io.py        | 30 +++++++++++++++++++++++++++++-
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/python/podio/frame.py b/python/podio/frame.py
index 3886fe28a..b955162ff 100644
--- a/python/podio/frame.py
+++ b/python/podio/frame.py
@@ -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()
 
diff --git a/python/podio/frame_iterator.py b/python/podio/frame_iterator.py
index d82ad2d70..6976ffb43 100644
--- a/python/podio/frame_iterator.py
+++ b/python/podio/frame_iterator.py
@@ -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:
@@ -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))
 
@@ -52,6 +58,8 @@ 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 '
@@ -59,7 +67,13 @@ def __getitem__(self, entry):
             '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
diff --git a/python/podio/root_io.py b/python/podio/root_io.py
index 0fcc75c83..e62347dd6 100644
--- a/python/podio/root_io.py
+++ b/python/podio/root_io.py
@@ -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
@@ -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).
 
@@ -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