Move the class to base_writer.py

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

@@ -1,7 +1,6 @@
 #!/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
@@ -10,29 +9,6 @@
 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:
-      try:
-        writer._writer.finish()  # pylint: disable=protected-access
-      except AttributeError:
-        pass
-
-
-_all_writers = AllWriters()
-atexit.register(_all_writers.finish)
-
-
 class Reader(BaseReaderMixin):
   """Reader class for reading podio root files."""
 
@@ -101,7 +77,7 @@ def __init__(self, filename):
         filename (str): The name of the output file
     """
     self._writer = podio.ROOTFrameWriter(filename)
-    _all_writers.add(self)
+    super().__init__()
 
 
 class RNTupleWriter(BaseWriterMixin):
@@ -113,4 +89,4 @@ def __init__(self, filename):
         filename (str): The name of the output file
     """
     self._writer = podio.ROOTNTupleWriter(filename)
-    _all_writers.add(self)
+    super().__init__()