refactor: Register headings with a Markdown processor rather than the on_page_content hook and the table of contents

Author: pawamoy

src/mkdocs_autorefs/__init__.py

@@ -14,6 +14,7 @@
     AutorefsExtension,
     AutorefsHookInterface,
     AutorefsInlineProcessor,
+    HeadingScannerTreeProcessor,
     fix_ref,
     fix_refs,
     relative_url,
@@ -31,6 +32,7 @@
     "Backlink",
     "BacklinkCrumb",
     "BacklinksTreeProcessor",
+    "HeadingScannerTreeProcessor",
     "fix_ref",
     "fix_refs",
     "relative_url",

src/mkdocs_autorefs/_internal/plugin.py

@@ -211,35 +211,32 @@ def on_page_markdown(self, markdown: str, page: Page, **kwargs: Any) -> str:  #
 
         Returns:
             The same Markdown. We only use this hook to keep a reference to the current page URL,
-                used during Markdown conversion by the anchor scanner tree processor.
+                used during Markdown conversion by the anchor/heading scanner tree processors.
         """
         # YORE: Bump 2: Remove line.
         self._url_to_page[page.url] = page
         self.current_page = page
         return markdown
 
     def on_page_content(self, html: str, page: Page, **kwargs: Any) -> str:  # noqa: ARG002
-        """Map anchors to URLs.
+        """Register breadcrumbs.
 
         Hook for the [`on_page_content` event](https://www.mkdocs.org/user-guide/plugins/#on_page_content).
-        In this hook, we map the IDs of every anchor found in the table of contents to the anchors absolute URLs.
-        This mapping will be used later to fix unresolved reference of the form `[title][identifier]` or
-        `[identifier][]`.
+        In this hook, we register breadcrumbs for each heading on each page.
+        These breadcrumbs are used to provide contextual information in backlinks.
 
         Arguments:
             html: HTML converted from Markdown.
             page: The related MkDocs page instance.
             kwargs: Additional arguments passed by MkDocs.
 
         Returns:
-            The same HTML. We only use this hook to map anchors to URLs.
+            The same HTML.
         """
         self.current_page = page
-        # Collect `std`-domain URLs.
-        if self.scan_toc:
-            _log.debug("Mapping identifiers to URLs for page %s", page.file.src_path)
+        if self.record_backlinks:
             for item in page.toc.items:
-                self.map_urls(page, item)
+                self._register_breadcrumbs(page, item)
         return html
 
     @event_priority(-50)  # Late, after mkdocstrings has finished loading inventories.
@@ -299,30 +296,43 @@ def on_env(self, env: Environment, /, *, config: MkDocsConfig, files: Files) ->
     # ----------------------------------------------------------------------- #
     # Utilities                                                               #
     # ----------------------------------------------------------------------- #
-    # TODO: Maybe stop exposing this method in the future.
+    # YORE: Bump 2: Remove block.
     def map_urls(self, page: Page, anchor: AnchorLink) -> None:
-        """Recurse on every anchor to map its ID to its absolute URL.
+        """Deprecated. Recurse on every anchor to map its ID to its absolute URL.
+
+        This method is deprecated and will be removed in a future release.
+        Please use the `register_anchor` or `register_url` methods instead.
 
         This method populates `self._primary_url_map` by side-effect.
 
         Arguments:
             page: The page containing the anchors.
             anchor: The anchor to process and to recurse on.
         """
-        return self._map_urls(page, anchor)
+        warn(
+            "The `map_urls` method is deprecated and will be removed in a future release. "
+            "Please use the `register_anchor` or `register_url` methods instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        self._map_urls(page, anchor)
 
-    def _map_urls(self, page: Page, anchor: AnchorLink, parent: BacklinkCrumb | None = None) -> None:
-        # YORE: Bump 2: Remove block.
+    # YORE: Bump 2: Remove block.
+    def _map_urls(self, page: Page | str, anchor: AnchorLink) -> None:
         if isinstance(page, str):
             try:
                 page = self._url_to_page[page]
             except KeyError:
                 page = self.current_page  # type: ignore[assignment]
 
-        self.register_anchor(page, anchor.id, title=anchor.title, primary=True)
+        self.register_anchor(page, anchor.id, title=anchor.title, primary=True)  # type: ignore[arg-type]
+        self._register_breadcrumbs(page, anchor)  # type: ignore[arg-type]
+
+    def _register_breadcrumbs(self, page: Page, anchor: AnchorLink, parent: BacklinkCrumb | None = None) -> None:
+        # Getting a breadcrumb has a side-effect of registering it in the breadcrumbs map.
         breadcrumb = self._get_breadcrumb(page, anchor, parent)
         for child in anchor.children:
-            self._map_urls(page, child, breadcrumb)
+            self._register_breadcrumbs(page, child, breadcrumb)
 
     def _get_breadcrumb(
         self,

src/mkdocs_autorefs/_internal/references.py

@@ -282,6 +282,45 @@ def _scan_anchors(self, parent: Element, pending_anchors: _PendingAnchors, last_
                 self.run(el)
 
 
+class HeadingScannerTreeProcessor(Treeprocessor):
+    """Tree processor to scan and register HTML headings."""
+
+    name: str = "mkdocs-autorefs-headings-scanner"
+    """The name of the tree processor."""
+
+    _htags: ClassVar[set[str]] = {"h1", "h2", "h3", "h4", "h5", "h6"}
+
+    def __init__(self, plugin: AutorefsPlugin, md: Markdown | None = None) -> None:
+        """Initialize the tree processor.
+
+        Parameters:
+            plugin: A reference to the autorefs plugin, to use its `register_anchor` method.
+        """
+        super().__init__(md)
+        self._plugin = plugin
+
+    def run(self, root: Element) -> None:
+        """Run the tree processor.
+
+        Arguments:
+            root: The root element of the tree.
+        """
+        if self._plugin.current_page is not None:
+            self._scan_headings(root)
+
+    def _scan_headings(self, parent: Element) -> None:
+        for el in parent:
+            if el.tag in self._htags:
+                if h_id := el.get("id"):
+                    self._plugin.register_anchor(
+                        self._plugin.current_page,  # type: ignore[arg-type]
+                        h_id,
+                        title=el.text,
+                    )
+            else:
+                self._scan_headings(el)
+
+
 class AutorefsExtension(Extension):
     """Markdown extension that transforms unresolved references into auto-references.
 
@@ -323,6 +362,13 @@ def extendMarkdown(self, md: Markdown) -> None:  # noqa: N802 (casing: parent me
             priority=168,  # Right after markdown.inlinepatterns.ReferenceInlineProcessor
         )
         if self.plugin is not None:
+            # Scan headings to register them.
+            if self.plugin.scan_toc:
+                md.treeprocessors.register(
+                    HeadingScannerTreeProcessor(self.plugin, md),
+                    HeadingScannerTreeProcessor.name,
+                    priority=0,
+                )
             # Markdown anchors require the `attr_list` extension.
             if self.plugin.scan_toc and "attr_list" in md.treeprocessors:
                 _log_enabling_markdown_anchors()

tests/test_backlinks.py

@@ -26,7 +26,10 @@ def test_get_backlinks() -> None:
     """Check that backlinks can be retrieved."""
     plugin = AutorefsPlugin()
     plugin.record_backlinks = True
-    plugin.map_urls(create_page("foo.html"), create_anchor_link("Foo", "foo"))
+    page = create_page("foo.html")
+    anchor = create_anchor_link("Foo", "foo")
+    plugin.register_anchor(page, anchor.id, title=anchor.title, primary=True)
+    plugin._register_breadcrumbs(page, anchor)
     plugin._primary_url_map["bar"] = ["bar.html#bar"]
     plugin._record_backlink("bar", "referenced-by", "foo", "foo.html")
     assert plugin.get_backlinks("bar", from_url="") == {

tests/test_references.py

@@ -344,18 +344,25 @@ def test_register_markdown_anchors() -> None:
     )
     assert plugin._primary_url_map == {
         "foo": ["page#heading-foo"],
+        "heading-foo": ["page#heading-foo"],
         "bar": ["page#bar"],
+        "heading-bar": ["page#heading-bar"],
         "alias1": ["page#heading-bar"],
         "alias2": ["page#heading-bar"],
         "alias3": ["page#alias3"],
         "alias4": ["page#heading-baz"],
+        "heading-baz": ["page#heading-baz"],
         "alias5": ["page#alias5"],
         "alias6": ["page#alias6"],
+        "heading-more1": ["page#heading-more1"],
         "alias7": ["page#alias7"],
         "alias8": ["page#alias8"],
         "alias9": ["page#heading-custom2"],
+        "heading-custom2": ["page#heading-custom2"],
         "alias10": ["page#alias10"],
         "aliasSame": ["page#same-heading-1", "page#same-heading-2"],
+        "same-heading-1": ["page#same-heading-1"],
+        "same-heading-2": ["page#same-heading-2"],
     }
 
 
@@ -380,6 +387,9 @@ def test_register_markdown_anchors_with_admonition() -> None:
         ),
     )
     assert plugin._primary_url_map == {
+        "heading-foo": ["page#heading-foo"],
+        "heading-bar": ["page#heading-bar"],
+        "heading-baz": ["page#heading-baz"],
         "alias1": ["page#alias1"],
         "alias2": ["page#heading-bar"],
         "alias3": ["page#alias3"],