bugfix for tracing reenablement in module stmt

ipyflow · Dec 5, 2023 · f83c14c · f83c14c
1 parent 9f3f60e
commit f83c14c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/core/ipyflow/tracing/ipyflow_tracer.py b/core/ipyflow/tracing/ipyflow_tracer.py
@@ -1436,8 +1436,9 @@ def before_stmt(self, _ret: None, stmt_id: int, frame: FrameType, *_, **__) -> N
             ):
                 self.after_stmt(None, prev_trace_stmt_in_cur_frame.stmt_id, frame)
         self.prev_trace_stmt_in_cur_frame = trace_stmt
-        if not self.is_tracing_enabled and self._try_reenable_tracing(
-            frame, dry_run=True
+        if not self.is_tracing_enabled and (
+            trace_stmt.is_module_stmt()
+            or self._try_reenable_tracing(frame, dry_run=True)
         ):
             self.after_stmt_reset_hook()
             self._tracked_enable_tracing()