fix: Support HTML toc labels with Python-Markdown 3.6+

pawamoy · pawamoy · commit c0d009000678 · 2024-04-02T20:40:19.000+02:00
diff --git a/src/mkdocstrings/extension.py b/src/mkdocstrings/extension.py
@@ -228,7 +228,7 @@ def _process_block(
         return rendered, handler, data
 
 
-class _PostProcessor(Treeprocessor):
+class _HeadingsPostProcessor(Treeprocessor):
     def run(self, root: Element) -> None:
         self._remove_duplicated_headings(root)
 
@@ -249,6 +249,17 @@ def _remove_duplicated_headings(self, parent: Element) -> None:
             parent.text = (parent.text or "") + carry_text
 
 
+class _TocLabelsTreeProcessor(Treeprocessor):
+    def run(self, root: Element) -> None:  # noqa: ARG002
+        self._override_toc_labels(self.md.toc_tokens)  # type: ignore[attr-defined]
+
+    def _override_toc_labels(self, tokens: list[dict[str, Any]]) -> None:
+        for token in tokens:
+            if (label := token.get("data-toc-label")) and token["name"] != label:
+                token["name"] = label
+                self._override_toc_labels(token["children"])
+
+
 class MkdocstringsExtension(Extension):
     """Our Markdown extension.
 
@@ -284,7 +295,12 @@ def extendMarkdown(self, md: Markdown) -> None:  # noqa: N802 (casing: parent me
             priority=75,  # Right before markdown.blockprocessors.HashHeaderProcessor
         )
         md.treeprocessors.register(
-            _PostProcessor(md),
-            "mkdocstrings_post",
+            _HeadingsPostProcessor(md),
+            "mkdocstrings_post_headings",
             priority=4,  # Right after 'toc'.
         )
+        # md.treeprocessors.register(
+        #     _TocLabelsTreeProcessor(md),
+        #     "mkdocstrings_post_toc_labels",
+        #     priority=4,  # Right after 'toc'.
+        # )
diff --git a/tests/test_extension.py b/tests/test_extension.py
@@ -105,26 +105,46 @@ def test_no_double_toc(ext_markdown: Markdown, expect_permalink: str) -> None:
         {
             "level": 1,
             "id": "aa",
+            "html": "aa",
             "name": "aa",
+            "data-toc-label": "",
             "children": [
                 {
                     "level": 2,
                     "id": "tests.fixtures.headings--foo",
+                    "html": "Foo",
                     "name": "Foo",
+                    "data-toc-label": "",
                     "children": [
                         {
                             "level": 4,
                             "id": "tests.fixtures.headings--bar",
+                            "html": "Bar",
                             "name": "Bar",
+                            "data-toc-label": "",
                             "children": [
-                                {"level": 6, "id": "tests.fixtures.headings--baz", "name": "Baz", "children": []},
+                                {
+                                    "level": 6,
+                                    "id": "tests.fixtures.headings--baz",
+                                    "html": "Baz",
+                                    "name": "Baz",
+                                    "data-toc-label": "",
+                                    "children": [],
+                                },
                             ],
                         },
                     ],
                 },
             ],
         },
-        {"level": 1, "id": "bb", "name": "bb", "children": []},
+        {
+            "level": 1,
+            "id": "bb",
+            "html": "bb",
+            "name": "bb",
+            "data-toc-label": "",
+            "children": [],
+        },
     ]
 
 
