diff --git a/python/podio/base_writer.py b/python/podio/base_writer.py index 6f2b5777a..5fab89bfa 100644 --- a/python/podio/base_writer.py +++ b/python/podio/base_writer.py @@ -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. @@ -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. diff --git a/python/podio/root_io.py b/python/podio/root_io.py index 6ebfbdac7..c4e0a14f2 100644 --- a/python/podio/root_io.py +++ b/python/podio/root_io.py @@ -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): @@ -88,3 +89,4 @@ def __init__(self, filename): filename (str): The name of the output file """ self._writer = podio.ROOTNTupleWriter(filename) + super().__init__() diff --git a/python/podio/sio_io.py b/python/podio/sio_io.py index a0d45236b..19e606683 100644 --- a/python/podio/sio_io.py +++ b/python/podio/sio_io.py @@ -56,3 +56,4 @@ def __init__(self, filename): filename (str): The name of the output file """ self._writer = podio.SIOFrameWriter(filename) + super().__init__()