From 93e55851dd68b8f08bc075e76355c43cb0a3aed7 Mon Sep 17 00:00:00 2001
From: Paulo Meira <10246101+PMeira@users.noreply.github.com>
Date: Mon, 29 Jan 2024 01:06:13 -0300
Subject: [PATCH] Events: tweak finalization and context references, avoiding
 calls to dead callbacks.

---
 dss_python_backend/events.py | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/dss_python_backend/events.py b/dss_python_backend/events.py
index 6e529c0..ad28ab8 100644
--- a/dss_python_backend/events.py
+++ b/dss_python_backend/events.py
@@ -1,3 +1,5 @@
+import atexit
+from weakref import WeakKeyDictionary
 from .enums import AltDSSEvent
 from . import ffi, lib
 
@@ -8,7 +10,7 @@
 )
 
 class EventCallbackManager:
-    _ctx_to_manager = {}
+    _ctx_to_manager = WeakKeyDictionary()
 
     def __init__(self, ctx):
         if ctx in EventCallbackManager._ctx_to_manager:
@@ -19,6 +21,23 @@ def __init__(self, ctx):
         for evt_type in AltDSSEvent:
             setattr(self, evt_type.name, [])
 
+
+    def unregister_all(self):
+        for evt_type in AltDSSEvent:
+            handlers = getattr(self, evt_type.name)
+            if not handlers:
+                continue
+
+            handlers[:] = []
+            lib.ctx_DSSEvents_UnregisterAlt(
+                self.ctx,
+                evt_type,
+                lib.altdss_python_util_callback
+            )
+
+    def __del__(self):
+        self.unregister_all()
+
     def register_func(self, evt: AltDSSEvent, func) -> bool:
         handlers = getattr(self, AltDSSEvent(evt).name)
         if len(handlers) == 0:
@@ -88,4 +107,14 @@ def altdss_python_util_callback(ctx, eventCode: int, step: int, ptr):
         lib.ctx_Error_Set_Description(ctx, f"Python callback exception: {ex}".encode())
 
 
+def _remove_callbacks():
+    '''
+    Remove all callbacks at exit. Since the native library may outlive the Python callbacks,
+    we need to remove the callbacks here to ensure they are not called.
+    '''
+    for ctx_mgr in EventCallbackManager._ctx_to_manager.values():
+        ctx_mgr.unregister_all()
+
+atexit.register(_remove_callbacks)
+
 __all__ = ['get_manager_for_ctx']