From c434f23750b47b4b070e1d3c0b1cc0be79e89015 Mon Sep 17 00:00:00 2001
From: Maor Kleinberger <kmaork@gmail.com>
Date: Wed, 17 Aug 2022 23:58:39 +0300
Subject: [PATCH] Fix done_callback loss

---
 madbg/debugger.py | 29 +++++++++++++++--------------
 setup.cfg         |  2 +-
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/madbg/debugger.py b/madbg/debugger.py
index c04ed39..af2bdde 100644
--- a/madbg/debugger.py
+++ b/madbg/debugger.py
@@ -45,8 +45,9 @@ def __init__(self, stdin, stdout, term_type):
         term_output = Vt100_Output.from_pty(stdout, term_type)
         super().__init__(pt_session_options=dict(input=term_input, output=term_output), stdin=stdin, stdout=stdout)
         self.use_rawinput = True
+        self.done_callback = None
 
-    def trace_dispatch(self, frame, event, arg, check_debugging_global=False, done_callback=None):
+    def trace_dispatch(self, frame, event, arg, check_debugging_global=False):
         """
         Overriding super to allow the check_debugging_global and done_callback args.
 
@@ -55,9 +56,9 @@ def trace_dispatch(self, frame, event, arg, check_debugging_global=False, done_c
         """
         if check_debugging_global:
             if self._DEBUGGING_GLOBAL in frame.f_globals:
-                self.set_trace(frame, done_callback=done_callback)
+                self.set_trace(frame)
             else:
-                return
+                return None
         bdb_quit = False
         try:
             return super().trace_dispatch(frame, event, arg)
@@ -65,22 +66,22 @@ def trace_dispatch(self, frame, event, arg, check_debugging_global=False, done_c
             bdb_quit = True
         finally:
             if self.quitting or bdb_quit:
-                # To debugger finalization
-                if done_callback is not None:
-                    done_callback()
+                self._on_done()
+
+    def _on_done(self):
+        if self.done_callback is not None:
+            self.done_callback()
+            self.done_callback = None
 
     def set_trace(self, frame=None, done_callback=None):
         """ Overriding super to add the done_callback argument, allowing cleanup after a debug session """
-        td = lambda *args: self.trace_dispatch(*args, done_callback=done_callback)
+        if done_callback is not None:
+            # set_trace was called again without the previous one exiting -
+            # happens on continue -> ctrl-c
+            self.done_callback = done_callback
         if frame is None:
             frame = currentframe().f_back
-        self.reset()
-        while frame:
-            frame.f_trace = td
-            self.botframe = frame
-            frame = frame.f_back
-        self.set_step()
-        sys.settrace(td)
+        return super().set_trace(frame)
 
     def do_continue(self, arg):
         """ Overriding super to add a print """
diff --git a/setup.cfg b/setup.cfg
index 791c381..d7ee46a 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -9,7 +9,7 @@ install_requires =
 
 [metadata]
 name = madbg
-version = 1.3.0
+version = 1.3.1
 description = A fully-featured remote debugger for python
 author = Maor Kleinberger
 author_email = kmaork@gmail.com