Add a class to clean up the writers

python/podio/root_io.py

@@ -1,13 +1,31 @@
 #!/usr/bin/env python3
 """Python module for reading root files containing podio Frames"""
 
+import atexit
 from ROOT import gSystem
 gSystem.Load('libpodioRootIO')  # noqa: E402
 from ROOT import podio  # noqa: E402 # pylint: disable=wrong-import-position
 
 from podio.base_reader import BaseReaderMixin  # pylint: disable=wrong-import-position
 from podio.base_writer import BaseWriterMixin  # pylint: disable=wrong-import-position
 
+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:
+      writer._writer.finish()
+
+_all_writers = AllWriters()
+atexit.register(_all_writers.finish)
 
 class Reader(BaseReaderMixin):
   """Reader class for reading podio root files."""
@@ -77,6 +95,7 @@ def __init__(self, filename):
         filename (str): The name of the output file
     """
     self._writer = podio.ROOTFrameWriter(filename)
+    _all_writers.add(self)
 
 
 class RNTupleWriter(BaseWriterMixin):
@@ -88,3 +107,4 @@ def __init__(self, filename):
         filename (str): The name of the output file
     """
     self._writer = podio.ROOTNTupleWriter(filename)
+    _all_writers.add(self)