fix: Display code that triggered an error in warning

sphinxext_altair/altairplot.py

@@ -233,6 +233,9 @@ def run(self) -> list[nodes.Element]:
         return result
 
 
+class AltairPlotWarning(UserWarning): ...
+
+
 def html_visit_altair_plot(self: altair_plot, node: nodes.Element) -> None:  # noqa: C901
     # Execute the code, saving output and namespace
     namespace = node["namespace"]
@@ -242,14 +245,19 @@ def html_visit_altair_plot(self: altair_plot, node: nodes.Element) -> None:  # n
             chart = eval_block(node["code"], namespace)
         stdout = f.getvalue()
     except Exception as err:
+        err_file = node["rst_source"]
+        line_no = node["rst_lineno"]
+        err_code = node["code"]
         msg = (
-            f"altair-plot: {node['rst_source']}:{node['rst_lineno']} "
-            f"Code Execution failed: {type(err).__name__}: {err!s}"
+            f"Code Execution failed.\n"
+            f"  {err_file}:{line_no}\n"
+            f"    {type(err).__name__}: {err!s}\n"
+            f"    {err_code}"
         )
         if node["strict"]:
             raise ValueError(msg) from err
         else:
-            warnings.warn(msg, stacklevel=1)
+            warnings.warn(msg, AltairPlotWarning, stacklevel=1)
             raise nodes.SkipNode from err
 
     if chart_name := node.get("chart-var-name", None):
@@ -298,11 +306,14 @@ def html_visit_altair_plot(self: altair_plot, node: nodes.Element) -> None:  # n
             )
             self.body.append(html)
         else:
+            err_file = node["rst_source"]
+            line_no = node["rst_lineno"]
             msg = (
-                f"altair-plot: {node['rst_source']}:{node['rst_lineno']} Malformed block. "
-                "Last line of code block should define a valid altair Chart object."
+                f"Malformed block.\n"
+                f"  {err_file}:{line_no}\n"
+                f"    Last line of code block should define a valid altair Chart object."
             )
-            warnings.warn(msg, stacklevel=1)
+            warnings.warn(msg, AltairPlotWarning, stacklevel=1)
         raise nodes.SkipNode
 
 

tests/conftest.py

@@ -10,5 +10,5 @@
 
 
 @pytest.fixture(scope="session")
-def rootdir():
+def rootdir() -> Path:
     return Path(__file__).parent / "roots"

tests/roots/test-altairplot/errors_warnings.rst

@@ -0,0 +1,8 @@
+Trigger a NameError
+---------------------------
+...
+
+.. altair-plot::
+
+    polars.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+

tests/test_altairplot.py

@@ -1,21 +1,32 @@
 # ruff: noqa: E501
 # Tests are inspired by the test suite of sphinx itself
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING, cast
+
 import pytest
 
 from altair import SCHEMA_URL
 from sphinxext_altair.altairplot import (
     VEGA_JS_URL_DEFAULT,
     VEGAEMBED_JS_URL_DEFAULT,
     VEGALITE_JS_URL_DEFAULT,
+    AltairPlotWarning,
     purge_altair_namespaces,
     validate_links,
 )
 
+if TYPE_CHECKING:
+    from sphinx.application import Sphinx
+
+    from sphinxext_altair.altairplot import BuildEnvironment
+
 
 @pytest.mark.parametrize("add_namespaces_attr", [True, False])
 @pytest.mark.sphinx(testroot="altairplot")
-def test_purge_altair_namespaces(add_namespaces_attr, app):
-    env = app.env
+def test_purge_altair_namespaces(add_namespaces_attr: bool, app: Sphinx) -> None:
+    env: BuildEnvironment = cast("BuildEnvironment", app.env)
     if add_namespaces_attr:
         env._altair_namespaces = {"docname": {}}
 
@@ -39,7 +50,7 @@ def test_purge_altair_namespaces(add_namespaces_attr, app):
         ("editor source", {"editor": True, "source": True, "export": False}),
     ],
 )
-def test_validate_links(links, expected):
+def test_validate_links(links: str, expected: str | bool | dict[str, bool]) -> None:
     if expected == "raise":
         with pytest.raises(
             ValueError, match=r"Following links are invalid: \['unknown'\]"
@@ -51,8 +62,15 @@ def test_validate_links(links, expected):
 
 
 @pytest.mark.sphinx(testroot="altairplot", freshenv=True)
-def test_altairplotdirective(app):
-    app.builder.build_all()
+def test_altairplotdirective(app: Sphinx) -> None:
+    with pytest.warns(
+        AltairPlotWarning,
+        match=re.compile(
+            r"errors_warnings\.rst:5\n.+polars\.DataFrame\(\{\"a\": \[1, 2, 3\], \"b\": \[4, 5, 6\]\}\)",
+            re.DOTALL,
+        ),
+    ):
+        app.builder.build_all()
     result = (app.outdir / "index.html").read_text(encoding="utf8")
     assert result.count("https://cdn.jsdelivr.net/npm/vega@") == 1
     assert result.count("https://cdn.jsdelivr.net/npm/vega-lite@") == 1