fix interpreter source traceback (#602)

When KIRIN_PYTHON_STACKTRACE is set to 0, the source traceback should
identify the statement in the original kernel that's causing the error.
This is currently blank due to a bug in the base interpreter error
handling; this PR fixes the issue.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: neelay893

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
     - id: isort
       name: isort (python)
 - repo: https://github.com/psf/black
-  rev: 25.1.0
+  rev: 26.1.0
   hooks:
     - id: black
 - repo: https://github.com/charliermarsh/ruff-pre-commit

pyproject.toml

@@ -73,7 +73,7 @@ include = ["src/kirin/*"]
 
 [tool.uv]
 dev-dependencies = [
-    "black>=24.10.0",
+    "black>=26.1.0",
     "coverage>=7.6.4",
     "ipykernel>=6.29.5",
     "ipython>=8.29.0",

src/kirin/dialects/math/_gen.py

@@ -60,9 +60,7 @@ def builtin_math_functions():
             ret_type = "types.Bool"
         else:
             ret_type = "types.Float"
-        f.write(
-            textwrap.dedent(
-                f"""
+        f.write(textwrap.dedent(f"""
 @statement(dialect=dialect)
 class {name}(ir.Statement):
     \"\"\"{name} statement, wrapping the math.{name} function
@@ -71,9 +69,7 @@ class {name}(ir.Statement):
     traits = frozenset({{ir.Pure(), lowering2.FromPythonCall()}})
 {fields}
     result: ir.ResultValue = info.result({ret_type})
-"""
-            )
-        )
+"""))
 
 
 with open(os.path.join(os.path.dirname(__file__), "interp.py"), "w") as f:
@@ -89,23 +85,19 @@ class {name}(ir.Statement):
         fields = ", ".join(
             [f"values[{idx}]" for idx, _ in enumerate(sig.parameters.keys())]
         )
-        implements.append(
-            f"""
+        implements.append(f"""
     @impl(stmts.{name})
     def {name}(self, interp, frame: Frame, stmt: stmts.{name}):
         values = frame.get_values(stmt.args)
-        return (math.{name}({fields}),)"""
-        )
+        return (math.{name}({fields}),)""")
 
     # Write the interpreter class
     methods = "\n\n".join(implements)
-    f.write(
-        f"""
+    f.write(f"""
 @dialect.register
 class MathMethodTable(MethodTable):
 {methods}
-"""
-    )
+""")
 
 # __init__.py
 with open(os.path.join(os.path.dirname(__file__), "__init__.py"), "w") as f:
@@ -124,14 +116,10 @@ class MathMethodTable(MethodTable):
             ret_type = "bool"
         else:
             ret_type = "float"
-        f.write(
-            textwrap.dedent(
-                f"""
+        f.write(textwrap.dedent(f"""
         @lowering2.wraps(stmts.{name})
         def {name}({", ".join(f"{arg}: {ret_type}" for arg in sig.parameters.keys())}) -> {ret_type}: ...
-        """
-            )
-        )
+        """))
     f.write("\n")
 
 for file in ["__init__.py", "interp.py", "stmts.py"]:
@@ -178,16 +166,12 @@ def {name}({", ".join(f"{arg}: {ret_type}" for arg in sig.parameters.keys())}) -
         args = ", ".join(arg for arg in sig.parameters.keys())
         inputs = ", ".join("0.42" for _ in sig.parameters.keys())
 
-        f.write(
-            textwrap.dedent(
-                f"""
+        f.write(textwrap.dedent(f"""
                 @basic
                 def {name}_func({args}):
                     return math.{name}({args})
 
                 def test_{name}():
                     truth = pymath.{name}({inputs})
                     assert ({name}_func({inputs}) - truth) < 1e-6
-                """
-            )
-        )
+                """))

src/kirin/interp/abc.py

@@ -264,7 +264,12 @@ def new_frame(
         self.state.push_frame(frame)
         try:
             yield frame
-        finally:
+        except BaseException:
+            # NOTE: Don't pop frames on exception so that eval_context
+            # can capture the full frame chain for the stack trace.
+            # The state is re-initialized on the next eval_context call.
+            raise
+        else:
             self.state.pop_frame()
 
     def frame_eval(

src/kirin/lowering/python/binding.py

@@ -14,8 +14,7 @@ class Binding(Generic[Params, RetType]):
 
     def __call__(self, *args: Params.args, **kwargs: Params.kwargs) -> RetType:
         raise NotImplementedError(
-            f"Binding of {self.parent.name} can \
-            only be called from a kernel"
+            f"Binding of {self.parent.name} can only be called from a kernel"
         )
 
 

test/lowering/test_assign.py

@@ -21,7 +21,7 @@ class MultiResult(ir.Statement):
 def test_multi_result():
     @dummy_dialect
     def multi_assign():
-        (x, y) = MultiResult()  # type: ignore
+        x, y = MultiResult()  # type: ignore
         return x, y
 
     stmt = multi_assign.callable_region.blocks[0].stmts.at(0)
@@ -33,7 +33,7 @@ def multi_assign():
 
         @dummy_dialect
         def multi_assign_error():
-            (x, y, z) = MultiResult()  # type: ignore
+            x, y, z = MultiResult()  # type: ignore
             return x, y, z