fix: Improve handling of errors within sessions

src/markdown_exec/formatters/base.py

@@ -87,7 +87,7 @@ def base_format(
         source_output = code
 
     try:
-        output = run(source_input, returncode=returncode, session=session, **extra)
+        output = run(source_input, returncode=returncode, session=session, id=id, **extra)
     except ExecutionError as error:
         identifier = id or extra.get("title", "")
         identifier = identifier and f"'{identifier}' "

src/markdown_exec/formatters/python.py

@@ -3,39 +3,65 @@
 from __future__ import annotations
 
 import traceback
+from collections import defaultdict
 from functools import partial
 from io import StringIO
 from typing import Any
 
 from markdown_exec.formatters.base import ExecutionError, base_format
 from markdown_exec.rendering import code_block
 
-_sessions: dict[str, dict] = {}
+_sessions_globals: dict[str, dict] = defaultdict(dict)
+_sessions_counter: dict[str | None, int] = defaultdict(int)
+_code_blocks: dict[str, list[str]] = {}
 
 
 def _buffer_print(buffer: StringIO, *texts: str, end: str = "\n", **kwargs: Any) -> None:  # noqa: ARG001
     buffer.write(" ".join(str(text) for text in texts) + end)
 
 
-def _run_python(code: str, session: str | None = None, **extra: str) -> str:
-    if session:
-        if session in _sessions:
-            exec_globals = _sessions[session]
-        else:
-            exec_globals = _sessions[session] = {}
+def _code_block_id(
+    id: str | None = None,  # noqa: A002
+    session: str | None = None,
+    title: str | None = None,
+) -> str:
+    _sessions_counter[session] += 1
+    if id:
+        code_block_id = f"id {id}"
+    elif session:
+        code_block_id = f"session {session}; n{_sessions_counter[session]}"
+        if title:
+            code_block_id = f"{code_block_id}; title {title}"
     else:
-        exec_globals = {}
+        code_block_id = f"n{_sessions_counter[session]}"
+        if title:
+            code_block_id = f"{code_block_id}; title {title}"
+    return f"<code block: {code_block_id}>"
+
+
+def _run_python(
+    code: str,
+    returncode: int | None = None,  # noqa: ARG001
+    session: str | None = None,
+    id: str | None = None,  # noqa: A002
+    **extra: str,
+) -> str:
+    title = extra.get("title", None)
+    code_block_id = _code_block_id(id, session, title)
+    _code_blocks[code_block_id] = code.split("\n")
+    exec_globals = _sessions_globals[session] if session else {}
+
     buffer = StringIO()
     exec_globals["print"] = partial(_buffer_print, buffer)
 
     try:
-        exec(code, exec_globals)  # noqa: S102
+        compiled = compile(code, filename=code_block_id, mode="exec")
+        exec(compiled, exec_globals)  # noqa: S102
     except Exception as error:  # noqa: BLE001
         trace = traceback.TracebackException.from_exception(error)
         for frame in trace.stack:
-            if frame.filename == "<string>":
-                frame.filename = "<executed code block>"
-                frame._line = code.split("\n")[frame.lineno - 1]  # type: ignore[attr-defined,operator]
+            if frame.filename.startswith("<code block: "):
+                frame._line = _code_blocks[frame.filename][frame.lineno - 1]  # type: ignore[attr-defined,operator]
         raise ExecutionError(code_block("python", "".join(trace.format()), **extra)) from error
     return buffer.getvalue()
 

tests/test_python.py

@@ -134,3 +134,31 @@ def test_sessions(md: Markdown) -> None:
     assert "b = 2" in html
     assert "ok" in html
     assert "ko" not in html
+
+
+def test_reporting_errors_in_sessions(md: Markdown, caplog: pytest.LogCaptureFixture) -> None:
+    """Assert errors and source lines are correctly reported across sessions.
+
+    Parameters:
+        md: A Markdown instance (fixture).
+        caplog: Pytest fixture to capture logs.
+    """
+    html = md.convert(
+        dedent(
+            """
+            ```python exec="1" session="a"
+            def fraise():
+                raise RuntimeError("strawberry")
+            ```
+
+            ```python exec="1" session="a"
+            print("hello")
+            fraise()
+            ```
+            """,
+        ),
+    )
+    assert "Traceback" in html
+    assert "strawberry" in html
+    assert "fraise()" in caplog.text
+    assert 'raise RuntimeError("strawberry")' in caplog.text