From 0492cf0875cd4c4cc8d1223f7b1036f8b009f778 Mon Sep 17 00:00:00 2001
From: Jan Kieseler <jkiesele@cern.ch>
Date: Fri, 13 Dec 2024 12:33:43 +0100
Subject: [PATCH] added better cleanup

---
 bind/minicalo_tools.py | 50 +++++++++++++++++++++++++++++-------------
 1 file changed, 35 insertions(+), 15 deletions(-)

diff --git a/bind/minicalo_tools.py b/bind/minicalo_tools.py
index c1524b5..64fee72 100644
--- a/bind/minicalo_tools.py
+++ b/bind/minicalo_tools.py
@@ -3,10 +3,15 @@
 import os
 import tempfile
 import pickle
+import weakref
+import atexit
 
 
 class TempFileManager:
-    def __init__(self, input_file=None, output_file=None, use_shm=True, file_permissions=0o644):
+
+    _instances = weakref.WeakSet()  # Track instances for cleanup on program exit
+
+    def __init__(self, input_file=None, output_file=None, use_shm=True, file_permissions=0o644, slave = False):
         """
         Initialize the TempFileManager.
         :param input_file: Path to the input file (optional).
@@ -17,6 +22,11 @@ class TempFileManager:
         self.output_file = output_file
         self.use_shm = use_shm and os.path.exists('/dev/shm')
         self.file_permissions = file_permissions
+        self.slave = slave
+
+        if not self.slave:
+            # Register this instance for cleanup
+            TempFileManager._instances.add(self)
 
     @staticmethod
     def slave_init(input_file, output_file):
@@ -26,7 +36,7 @@ class TempFileManager:
         :param output_file: Path to the output file.
         :return: A TempFileManager instance with the provided filenames.
         """
-        return TempFileManager(input_file=input_file, output_file=output_file)
+        return TempFileManager(input_file=input_file, output_file=output_file, slave=True)
 
     def __enter__(self):
         """
@@ -74,16 +84,26 @@ class TempFileManager:
         with open(self.output_file, 'rb') as f:
             return pickle.load(f)
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def cleanup(self):
         """
         Cleanup temporary files when exiting the context.
         """
-        if os.path.exists(self.input_file):
-            os.remove(self.input_file)
-        if os.path.exists(self.output_file):
-            os.remove(self.output_file)
+        for file_path in [self.input_file, self.output_file]:
+            if file_path and os.path.exists(file_path):
+                try:
+                    os.remove(file_path)
+                except OSError:
+                    pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.cleanup()
 
 
+# Ensure all instances are cleaned up on exit
+@atexit.register
+def cleanup_all_temp_files():
+    for instance in TempFileManager._instances:
+        instance.cleanup()
 
 
 import unittest
@@ -115,8 +135,8 @@ class TestTempFileManager(unittest.TestCase):
             tfm.dump_input(obj)
             
             # Read it back to ensure correctness
-            with open(tfm.input_file, 'r') as f:
-                data = json.load(f)
+            with open(tfm.input_file, 'rb') as f:
+                data = pickle.load(f)
             self.assertEqual(data, obj)
 
     def test_slave_init(self):
@@ -127,8 +147,8 @@ class TestTempFileManager(unittest.TestCase):
         try:
             # Create test input file
             obj = {"key": "test"}
-            with open(input_file, 'w') as f:
-                json.dump(obj, f)
+            with open(input_file, 'wb') as f:
+                pickle.dump(obj, f)
             
             # Initialize TempFileManager with slave_init
             tfm = TempFileManager.slave_init(input_file, output_file)
@@ -140,8 +160,8 @@ class TestTempFileManager(unittest.TestCase):
             # Write to the output file and verify
             updated_obj = {"key": "updated"}
             tfm.dump_output(updated_obj)
-            with open(output_file, 'r') as f:
-                output_data = json.load(f)
+            with open(output_file, 'rb') as f:
+                output_data = pickle.load(f)
             self.assertEqual(output_data, updated_obj)
         finally:
             # Cleanup test files
@@ -173,5 +193,5 @@ class TestTempFileManager(unittest.TestCase):
             permissions = stat.S_IMODE(os.stat(tfm.input_file).st_mode)
             self.assertEqual(permissions, 0o600)
 
-#if __name__ == "__main__":
-#    unittest.main()
+if __name__ == "__main__":
+    unittest.main()