Fix crash when a writer is in the global namespace

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.

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__()

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__()