Fix traceback

.github/workflows/main.yml

@@ -20,7 +20,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python-version: [ '3.9', '3.10', '3.11', '3.12', '3.13' ]
+        python-version: [ '3.10', '3.10', '3.11', '3.12', '3.13', '3.14' ]
 
     steps:
     - name: Checkout

pyproject.toml

@@ -8,16 +8,16 @@ version = "0.3.4"
 description = "An asynchronous Python Jupyter kernel"
 readme = "README.md"
 license = {text = "MIT"}
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 authors = [{name = "David Brochart", email = "david.brochart@gmail.com"}]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Programming Language :: Python",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]

src/akernel/code.py

@@ -149,10 +149,10 @@ def get_async_code(self) -> str:
         gtree = self.get_async_ast()
         return ast.unparse(gtree)
 
-    def get_async_bytecode(self) -> CodeType:
+    def get_async_bytecode(self, task_i: int) -> CodeType:
         tree = self.get_async_ast()
         #tree = gast.gast_to_ast(gtree)
-        bytecode = compile(tree, filename="<string>", mode="exec")
+        bytecode = compile(tree, filename=f"<cell-{task_i}>", mode="exec")
         return bytecode
 
     def make_react(self):

src/akernel/execution.py

@@ -25,7 +25,7 @@ def pre_execute(
 
     try:
         transform = Transform(code, task_i, react)
-        async_bytecode = transform.get_async_bytecode()
+        async_bytecode = transform.get_async_bytecode(task_i)
         exec(async_bytecode, globals_, locals_)
     except SyntaxError as e:
         exception = e

src/akernel/kernel.py

@@ -39,6 +39,7 @@ class Kernel:
     kernel_mode: str
     cell_done: Dict[int, Event]
     running_cells: Dict[int, asyncio.Task]
+    _source_map: Dict[str, str]
     task_i: int
     execution_count: int
     execution_state: str
@@ -85,6 +86,7 @@ def __init__(
         self.globals = {}
         self.locals = {}
         self._chain_execution = not self.concurrent_kernel
+        self._source_map = {}
         self.cell_done = {}
         self.running_cells = {}
         self.task_i = 0
@@ -374,13 +376,14 @@ async def execute_and_finish(
         parent_header = parent["header"]
         traceback, exception = [], None
         namespace = self.get_namespace(parent_header)
+        self._source_map[f"<cell-{task_i}>"] = code
         try:
             result = await self.locals[namespace][f"__async_cell{task_i}__"]()
         except KeyboardInterrupt:
             self.interrupt()
         except Exception as e:
             exception = e
-            traceback = get_traceback(code, e, execution_count)
+            traceback = get_traceback(code, e, execution_count, self._source_map)
         else:
             await self.show_result(result, self.globals[namespace], parent_header)
             cache_execution(self.cache, cache_info, self.globals[namespace], result)

src/akernel/traceback.py

@@ -2,12 +2,12 @@
 
 import sys
 import types
-from typing import cast
+from typing import Dict, Optional, cast
 
 from colorama import Fore, Style  # type: ignore
 
 
-def get_traceback(code: str, exception, execution_count: int = 0):
+def get_traceback(code: str, exception, execution_count: int = 0, source_map: Optional[Dict[str, str]] = None):
     exc_info = sys.exc_info()
     tb = cast(types.TracebackType, exc_info[2])
     while True:
@@ -26,22 +26,30 @@ def get_traceback(code: str, exception, execution_count: int = 0):
     traceback = ["Traceback (most recent call last):"]
     for frame in stack:
         filename = frame.f_code.co_filename
-        if filename == "<string>":
-            filename = f"{Fore.CYAN}Cell{Style.RESET_ALL} {Fore.GREEN}{execution_count}"
-            f"{Style.RESET_ALL}"
+        if filename.startswith("<cell-"):
+            source = source_map.get(filename, "") if source_map else ""
+            display_filename = (
+                f"{Fore.CYAN}Cell{Style.RESET_ALL} {Fore.GREEN}{int(filename[6:-1]) + 1}"
+                f"{Style.RESET_ALL}"
+            )
         else:
             with open(filename) as f:
-                code = f.read()
-            filename = f"{Fore.CYAN}File{Style.RESET_ALL} {Fore.GREEN}{filename}{Style.RESET_ALL}"
-        if frame.f_code.co_name.startswith("__async_cell"):
-            name = "<module>"
-        else:
-            name = frame.f_code.co_name
+                source = f.read()
+            display_filename = (
+                f"{Fore.CYAN}File{Style.RESET_ALL} {Fore.GREEN}{filename}"
+                f"{Style.RESET_ALL}"
+            )
+        name = "<module>" if frame.f_code.co_name.startswith("__async_cell") else frame.f_code.co_name
         trace = [
-            f"{filename} in {Fore.CYAN}{name}{Style.RESET_ALL}, {Fore.CYAN}line{Style.RESET_ALL} "
-            f"{Fore.GREEN}{frame.f_lineno}{Style.RESET_ALL}:"
+            f"  {display_filename}, "
+            f"{Fore.CYAN}line{Style.RESET_ALL} "
+            f"{Fore.GREEN}{frame.f_lineno}{Style.RESET_ALL}, "
+            f"in {Fore.CYAN}{name}{Style.RESET_ALL}:"
         ]
-        trace.append(code.splitlines()[frame.f_lineno - 1])
+        if source:
+            trace.append("    " + source.splitlines()[frame.f_lineno - 1].lstrip())
         traceback += trace
-    traceback += [f"{Fore.RED}{type(exception).__name__}{Style.RESET_ALL}: {exception.args[0]}"]
+    traceback += [
+        f"{Fore.RED}{type(exception).__name__}{Style.RESET_ALL}: {exception.args[0]}"
+    ]
     return traceback