Skip to content

Commit

Permalink
Fix crash when a writer is in the global namespace (#539)
Browse files Browse the repository at this point in the history
* Add a class to clean up the writers
  • Loading branch information
jmcarcell authored Jan 16, 2024
1 parent 532234f commit 276cdaa
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 0 deletions.
29 changes: 29 additions & 0 deletions python/podio/base_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,31 @@
"""Python module for defining the basic writer interface that is used by the
backend specific bindings"""

import atexit


class AllWriters:
"""Class to manage all writers in the program
so that they can be properly finished at the end of the program
"""
writers = []

def add(self, writer):
"""Add a writer to the list of managed writers"""
self.writers.append(writer)

def finish(self):
"""Finish all managed writers"""
for writer in self.writers:
try:
writer._writer.finish() # pylint: disable=protected-access
except AttributeError:
pass


_all_writers = AllWriters()
atexit.register(_all_writers.finish)


class BaseWriterMixin:
"""Mixin class that defines the base interface of the writers.
Expand All @@ -11,6 +36,10 @@ class BaseWriterMixin:
- _writer: The actual writer that is able to write frames
"""

def __init__(self):
"""Initialize the writer"""
_all_writers.add(self)

def write_frame(self, frame, category, collections=None):
"""Write the given frame under the passed category, optionally limiting the
collections that are written.
Expand Down
2 changes: 2 additions & 0 deletions python/podio/root_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ def __init__(self, filename):
filename (str): The name of the output file
"""
self._writer = podio.ROOTFrameWriter(filename)
super().__init__()


class RNTupleWriter(BaseWriterMixin):
Expand All @@ -88,3 +89,4 @@ def __init__(self, filename):
filename (str): The name of the output file
"""
self._writer = podio.ROOTNTupleWriter(filename)
super().__init__()
1 change: 1 addition & 0 deletions python/podio/sio_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,3 +56,4 @@ def __init__(self, filename):
filename (str): The name of the output file
"""
self._writer = podio.SIOFrameWriter(filename)
super().__init__()

0 comments on commit 276cdaa

Please sign in to comment.