From d674b7e94250cbc70e6c9f290a6f910fd64a01d4 Mon Sep 17 00:00:00 2001
From: Talley Lambert <talley.lambert@gmail.com>
Date: Thu, 28 Mar 2024 15:12:14 -0400
Subject: [PATCH] improve stack inspection

---
 src/psygnal/_exceptions.py | 56 ++++++++++++++++++++++++++------------
 src/psygnal/_signal.py     | 14 +++++-----
 2 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/src/psygnal/_exceptions.py b/src/psygnal/_exceptions.py
index efb7b61a..eaca6dcc 100644
--- a/src/psygnal/_exceptions.py
+++ b/src/psygnal/_exceptions.py
@@ -1,11 +1,19 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import inspect
+from contextlib import suppress
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import psygnal
 
 if TYPE_CHECKING:
     from ._signal import SignalInstance
 
 
+ROOT = str(Path(psygnal.__file__).parent)
+
+
 class EmitLoopError(Exception):
     """Error type raised when an exception occurs during a callback."""
 
@@ -16,10 +24,11 @@ def __init__(
         exc: BaseException,
         signal: SignalInstance | None = None,
     ) -> None:
-        self.__cause__ = exc  # mypyc doesn't set this, but uncompiled code would
+        self.__cause__ = exc
 
+        # grab the signal name or repr
         if signal is None:  # pragma: no cover
-            sig_name = ""
+            sig_name: Any = ""
         else:
             if instsance := signal.instance:
                 inst_class = instsance.__class__
@@ -28,20 +37,33 @@ def __init__(
                     mod += "."
                 sig_name = f"{mod}{inst_class.__qualname__}.{signal.name}"
             else:
-                sig_name = repr(signal)
+                sig_name = signal
 
-        msg = f"\n\nWhile emitting signal {sig_name!r}, an error occurred in a callback"
+        etype = exc.__class__.__name__  # name of the exception raised by callback.
+        msg = (
+            f"\n\nWhile emitting signal {sig_name!r}, a {etype} occurred in a callback"
+        )
         if tb := exc.__traceback__:
-            while tb and tb.tb_next is not None:
-                tb = tb.tb_next
-            frame = tb.tb_frame
-            filename = frame.f_code.co_filename
-            func_name = getattr(frame.f_code, "co_qualname", frame.f_code.co_name)
-            msg += f":\n  File {filename}:{frame.f_lineno}, in {func_name}\n"
-            if frame.f_locals:
-                msg += "  With local variables:\n"
-                for name, value in frame.f_locals.items():
-                    msg += f"    {name} = {value!r}\n"
-
-        msg += f"\nSee {exc.__class__.__name__} above for details."
+            msg += ":\n"
+
+            # get the first frame in the stack that is not in the psygnal package
+            with suppress(Exception):
+                fi = next(fi for fi in inspect.stack() if ROOT not in fi.filename)
+                msg += f"\n  Signal emitted at: {fi.filename}:{fi.lineno}, in {fi.function}\n"  # noqa: E501
+                if fi.code_context:
+                    msg += f"    >  {fi.code_context[0].strip()}\n"
+
+            # get the last frame in the traceback, the one that raised the exception
+            with suppress(Exception):
+                fi = inspect.getinnerframes(tb)[-1]
+                msg += f"\n  Callback error at: {fi.filename}:{fi.lineno}, in {fi.function}\n"  # noqa: E501
+                if fi.code_context:
+                    msg += f"    >  {fi.code_context[0].strip()}\n"
+                if flocals := fi.frame.f_locals:
+                    msg += "    Local variables:\n"
+                    for name, value in flocals.items():
+                        if name not in ("self", "cls"):
+                            msg += f"       {name} = {value!r}\n"
+
+        msg += f"\nSee {etype} above for original traceback."
         super().__init__(msg)
diff --git a/src/psygnal/_signal.py b/src/psygnal/_signal.py
index 5e393998..48310a33 100644
--- a/src/psygnal/_signal.py
+++ b/src/psygnal/_signal.py
@@ -1179,11 +1179,7 @@ def __call__(
         self, *args: Any, check_nargs: bool = False, check_types: bool = False
     ) -> None:
         """Alias for `emit()`. But prefer using `emit()` for clarity."""
-        return self.emit(
-            *args,
-            check_nargs=check_nargs,
-            check_types=check_types,
-        )
+        return self.emit(*args, check_nargs=check_nargs, check_types=check_types)
 
     def _run_emit_loop(self, args: tuple[Any, ...]) -> None:
         with self._lock:
@@ -1203,8 +1199,12 @@ def _run_emit_loop(self, args: tuple[Any, ...]) -> None:
                     f"RecursionError when "
                     f"emitting signal {self.name!r} with args {args}"
                 ) from e
-            except Exception as e:
-                raise EmitLoopError(exc=e, signal=self) from e
+            except Exception as cb_err:
+                loop_err = EmitLoopError(exc=cb_err, signal=self).with_traceback(
+                    cb_err.__traceback__
+                )
+                # this comment will show up in the traceback
+                raise loop_err from cb_err  # emit() call ABOVE || callback error BELOW
             finally:
                 self._recursion_depth -= 1
                 # we're back to the root level of the emit loop, reset max_depth