fix: Preserve ToC labels for stashed headings

src/markdown_exec/_internal/processors.py

@@ -5,6 +5,7 @@
 
 import copy
 import re
+from html.parser import HTMLParser
 from typing import TYPE_CHECKING
 from xml.etree.ElementTree import Element
 
@@ -16,6 +17,41 @@
     from markupsafe import Markup
 
 
+class _HeadingLabelsParser(HTMLParser):
+    """Extract visible heading labels from rendered HTML."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.labels: dict[str, str] = {}
+        self._current_id: str | None = None
+        self._current_label: list[str] | None = None
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        if HeadingReportingTreeprocessor.regex.fullmatch(tag):
+            self._current_id = dict(attrs).get("id")
+            self._current_label = []
+
+    def handle_data(self, data: str) -> None:
+        if self._current_label is not None:
+            self._current_label.append(data)
+
+    def handle_endtag(self, tag: str) -> None:
+        if (
+            self._current_label is not None
+            and HeadingReportingTreeprocessor.regex.fullmatch(tag)
+        ):
+            if self._current_id is not None:
+                self.labels[self._current_id] = "".join(self._current_label)
+            self._current_id = None
+            self._current_label = None
+
+
+def _heading_labels(markup: str) -> dict[str, str]:
+    parser = _HeadingLabelsParser()
+    parser.feed(markup)
+    return parser.labels
+
+
 # code taken from mkdocstrings, credits to @oprypin
 class IdPrependingTreeprocessor(Treeprocessor):
     """Prepend the configured prefix to IDs of all HTML elements."""
@@ -72,7 +108,11 @@ def run(self, root: Element) -> None:
                 el = copy.copy(el)  # noqa: PLW2901
                 # 'toc' extension's first pass (which we require to build heading stubs/ids) also edits the HTML.
                 # Undo the permalink edit so we can pass this heading to the outer pass of the 'toc' extension.
-                if len(el) > 0 and el[-1].get("class") == self.md.treeprocessors["toc"].permalink_class:  # type: ignore[attr-defined]
+                if (
+                    len(el) > 0
+                    and el[-1].get("class")
+                    == self.md.treeprocessors["toc"].permalink_class
+                ):  # type: ignore[attr-defined]
                     del el[-1]
                 self.headings.append(el)
 
@@ -98,11 +138,36 @@ def run(self, root: Element) -> None:
         if not self.headings:
             return
 
+        raw_html_blocks = self.md.htmlStash.rawHtmlBlocks
+
         for el in root.iter():
             match = HTML_PLACEHOLDER_RE.match(el.text or "")
             if match:
                 counter = int(match.group(1))
-                markup: Markup = self.md.htmlStash.rawHtmlBlocks[counter]  # type: ignore[assignment]
+                if counter < len(raw_html_blocks):
+                    # Other extensions can duplicate headings next to a stashed
+                    # HTML block so the ToC extension can see them. If their
+                    # labels contain inner stash placeholders, derive labels
+                    # from the actual rendered HTML instead.
+                    heading_labels = _heading_labels(str(raw_html_blocks[counter]))
+                    duplicated_headings = [
+                        child
+                        for child in el
+                        if HeadingReportingTreeprocessor.regex.fullmatch(child.tag)
+                    ]
+                    for heading in duplicated_headings:
+                        if label := heading_labels.get(heading.attrib.get("id", "")):
+                            heading.set("data-toc-label", label)
+
+            if HeadingReportingTreeprocessor.regex.fullmatch(el.tag):
+                continue
+
+            if match:
+                counter = int(match.group(1))
+                if counter >= len(raw_html_blocks):
+                    continue
+
+                markup: Markup = raw_html_blocks[counter]  # type: ignore[assignment]
                 if headings := self.headings.get(markup):
                     div = Element("div", {"class": "markdown-exec"})
                     div.extend(headings)
@@ -121,7 +186,8 @@ def run(self, root: Element) -> None:
 
     def _remove_duplicated_headings(self, parent: Element) -> None:
         carry_text = ""
-        for el in reversed(parent):  # Reversed mainly for the ability to mutate during iteration.
+        # Reversed mainly for the ability to mutate during iteration.
+        for el in reversed(parent):
             if el.tag == "div" and el.get("class") == "markdown-exec":
                 # Delete the duplicated headings along with their container, but keep the text (i.e. the actual HTML).
                 carry_text = (el.text or "") + carry_text

tests/test_headings.py

@@ -2,10 +2,15 @@
 
 from __future__ import annotations
 
+import os
+import subprocess
+import sys
 from textwrap import dedent
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from markdown import Markdown
 
 
@@ -28,3 +33,95 @@ def test_headings_removal(md: Markdown) -> None:
         ),
     )
     assert 'class="markdown-exec"' not in html
+
+
+def test_mkdocstrings_heading_placeholders_do_not_leak_into_toc(tmp_path: Path) -> None:
+    """Mkdocstrings headings with stashed inline code should not corrupt the ToC."""
+    (tmp_path / "docs").mkdir()
+    package = tmp_path / "repro_pkg"
+    package.mkdir()
+    (package / "__init__.py").write_text("", encoding="utf8")
+    (package / "repro97.py").write_text(
+        dedent(
+            '''\
+            """Summary.
+
+            # Heading `1`
+
+            Text with `code`.
+
+            # Heading `2`
+
+            Text with `code again`.
+
+            # Heading `3`
+
+            Text with `code again and again`.
+
+            # Heading `4`.
+
+            Text with `code again and again again`.
+            """
+            ''',
+        ),
+        encoding="utf8",
+    )
+    (tmp_path / "docs" / "index.md").write_text(
+        dedent(
+            """\
+            # Executed block
+
+            ```python exec="true" source="material-block"
+            print("hello from the exec block")
+            ```
+
+            # Documented module
+
+            ::: repro_pkg.repro97
+            """,
+        ),
+        encoding="utf8",
+    )
+    (tmp_path / "mkdocs.yml").write_text(
+        dedent(
+            """\
+            site_name: repro97
+            plugins:
+              - search
+              - markdown-exec
+              - mkdocstrings:
+                  handlers:
+                    python:
+                      paths: [.]
+            markdown_extensions:
+              - toc
+              - pymdownx.superfences
+              - pymdownx.highlight
+              - pymdownx.inlinehilite
+            """,
+        ),
+        encoding="utf8",
+    )
+    env = os.environ | {
+        "PYTHONPATH": os.pathsep.join(
+            path for path in (str(tmp_path), os.environ.get("PYTHONPATH", "")) if path
+        ),
+    }
+
+    result = subprocess.run(
+        [sys.executable, "-m", "mkdocs", "build", "--strict", "-q"],
+        cwd=tmp_path,
+        env=env,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+    assert result.returncode == 0, result.stderr
+    html = (tmp_path / "site" / "index.html").read_text(encoding="utf8")
+    assert "wzxhzdk" not in html
+    assert "Heading 1</a>" in html
+    assert "Heading 2</a>" in html
+    assert "Heading 3</a>" in html
+    assert "Heading 4.</a>" in html
+    assert "Heading print(&quot;hello from the exec block&quot;)" not in html