performance: iterate over raw HTML blocks only once per post-processor

Signed-off-by: Timothée Mazzucotelli <dev@pawamoy.fr>

python/tests/unit/extensions/test_glightbox.py

@@ -30,6 +30,8 @@ import pytest
 from tests.unit.extensions.conftest import soup
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+
     from markdown import Markdown
 
 # ---------------------------------------------------------------------------
@@ -48,6 +50,24 @@ def _glightbox(**kwargs: object) -> dict[str, Any]:
     }
 
 
+class _TrackingBlocks(list[str]):
+    """Track indexes read from stashed blocks."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.visited: list[int] = []
+
+    def __iter__(self) -> Iterator[str]:
+        for index, value in enumerate(super().__iter__()):
+            self.visited.append(index)
+            yield value
+
+    def __getitem__(self, index: Any) -> Any:
+        if isinstance(index, int):
+            self.visited.append(index)
+        return super().__getitem__(index)
+
+
 # ---------------------------------------------------------------------------
 # Basic wrapping
 # ---------------------------------------------------------------------------
@@ -388,3 +408,34 @@ class TestPostprocessor:
     def test_raw_html_skip_class_not_wrapped(self, md: Markdown) -> None:
         html = soup(md.convert('<img src="raw.png" class="off-glb" />'))
         assert html.select_one("a.glightbox") is None
+
+    @pytest.mark.parametrize(
+        "md",
+        [pytest.param(_glightbox(), id="default")],
+        indirect=["md"],
+    )
+    def test_toc_does_not_rescan_stash(
+        self, md: Markdown, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """TOC rendering does not cause processed blocks to be scanned again."""
+        blocks = _TrackingBlocks()
+        md.htmlStash.rawHtmlBlocks = blocks
+        processor = md.postprocessors["glightbox"]
+        original = processor.run
+        visits: list[list[int]] = []
+
+        def run(text: str) -> str:
+            start = len(blocks.visited)
+            text = original(text)
+            visits.append(blocks.visited[start:])
+            return text
+
+        monkeypatch.setattr(processor, "run", run)
+        html = soup(
+            md.convert('# One\n\n## Two\n\n## Three\n\n<img src="raw.png" />')
+        )
+
+        assert len(visits) > 1
+        assert visits[0] == list(range(len(blocks)))
+        assert all(not visited for visited in visits[1:])
+        assert len(html.select("a.glightbox")) == 1

python/tests/unit/extensions/test_links.py

@@ -23,16 +23,56 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import pytest
 from markdown import Markdown
 
 from zensical.extensions.links import (
     LinksExtension,
+    LinksPostprocessor,
     _is_relative,
     _md_path_to_html,
     _rewrite_url,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from typing import Any
+
+
+class _TrackingBlocks(list[str]):
+    """Track indexes read from stashed blocks."""
+
+    def __init__(self, values: list[str]) -> None:
+        super().__init__(values)
+        self.visited: list[int] = []
+
+    def __iter__(self) -> Iterator[str]:
+        for index, value in enumerate(super().__iter__()):
+            self.visited.append(index)
+            yield value
+
+    def __getitem__(self, index: Any) -> Any:
+        if isinstance(index, int):
+            self.visited.append(index)
+        return super().__getitem__(index)
+
+
+class _TrackingPostprocessor(LinksPostprocessor):
+    """Record which stash indexes each invocation reads."""
+
+    def __init__(self, md: Markdown, blocks: _TrackingBlocks) -> None:
+        super().__init__(md, "guide/page.md", True)
+        self._blocks = blocks
+        self.visits: list[list[int]] = []
+
+    def run(self, text: str) -> str:
+        start = len(self._blocks.visited)
+        text = super().run(text)
+        self.visits.append(self._blocks.visited[start:])
+        return text
+
 
 @pytest.mark.parametrize(
     ("path", "directory_urls", "expected"),
@@ -109,3 +149,21 @@ def test_rewrites_links_in_stashed_raw_html() -> None:
     assert md.convert('<div><a href="guide.md">Guide</a></div>') == (
         '<div><a href="guide/">Guide</a></div>'
     )
+
+
+def test_postprocessor_does_not_rescan_stash_for_toc() -> None:
+    """TOC rendering does not cause processed blocks to be scanned again."""
+    md = Markdown(extensions=["toc"])
+    blocks = _TrackingBlocks([])
+    md.htmlStash.rawHtmlBlocks = blocks
+    processor = _TrackingPostprocessor(md, blocks)
+    md.postprocessors.register(processor, processor.name, 29)
+
+    html = md.convert(
+        '# One\n\n## Two\n\n## Three\n\n<a href="other.md">other</a>'
+    )
+
+    assert len(processor.visits) > 1
+    assert processor.visits[0] == list(range(len(blocks)))
+    assert all(not visited for visited in processor.visits[1:])
+    assert '<a href="../other/">other</a>' in html

python/zensical/extensions/glightbox.py

@@ -216,6 +216,12 @@ class GlightboxTreeprocessor(Treeprocessor, ProcessorMixin):
         )
 
 
+# Python-Markdown's `toc` treeprocessor invokes every postprocessor while
+# rendering each heading and the generated table of contents, before Markdown
+# invokes them again for the complete document. Thus, this processor can run
+# several times against the same HTML stash. Since the stash grows by appending
+# blocks, a cursor both prevents wrapping images more than once and avoids
+# repeatedly scanning the already processed prefix.
 class GlightboxPostprocessor(Postprocessor, ProcessorMixin):
     """Wraps stashed images in anchors, delegating to the treeprocessor.
 
@@ -233,16 +239,16 @@ class GlightboxPostprocessor(Postprocessor, ProcessorMixin):
         self._skip_classes = self.SKIP_CLASSES | frozenset(
             self.config.skip_classes
         )
-        self._processed: set[int] = set()
+        self._cursor = 0
 
     def run(self, text: str) -> str:
         """Wrap images in stashed HTML blocks."""
-        for i, raw in enumerate(self.md.htmlStash.rawHtmlBlocks):
-            if i not in self._processed:
-                self.md.htmlStash.rawHtmlBlocks[i] = _RE.sub(
-                    self._maybe_process, raw
-                )
-                self._processed.add(i)
+        blocks = self.md.htmlStash.rawHtmlBlocks
+        while self._cursor < len(blocks):
+            blocks[self._cursor] = _RE.sub(
+                self._maybe_process, blocks[self._cursor]
+            )
+            self._cursor += 1
 
         # Return text unmodified, as we only need to modify the stashed raw HTML
         # blocks, which will later be reinstated by the raw HTML postprocessor

python/zensical/extensions/links.py

@@ -85,6 +85,12 @@ class LinksTreeprocessor(Treeprocessor):
                 el.set(key, url)
 
 
+# Python-Markdown's `toc` treeprocessor invokes every postprocessor while
+# rendering each heading and the generated table of contents, before Markdown
+# invokes them again for the complete document. Thus, this processor can run
+# several times against the same HTML stash. Since the stash grows by appending
+# blocks, a cursor both prevents rewriting URLs more than once and avoids
+# repeatedly scanning the already processed prefix.
 class LinksPostprocessor(Postprocessor):
     """Rewrites relative links in stashed raw HTML blocks.
 
@@ -103,16 +109,16 @@ class LinksPostprocessor(Postprocessor):
         super().__init__(md)
         self._path = path
         self._use_directory_urls = use_directory_urls
-        self._processed: set[int] = set()
+        self._cursor = 0
 
     def run(self, text: str) -> str:
         """Rewrite `href` and `src` attributes of stashed HTML blocks."""
-        for i, raw in enumerate(self.md.htmlStash.rawHtmlBlocks):
-            if i not in self._processed:
-                self.md.htmlStash.rawHtmlBlocks[i] = _RE.sub(
-                    self._maybe_process, raw
-                )
-                self._processed.add(i)
+        blocks = self.md.htmlStash.rawHtmlBlocks
+        while self._cursor < len(blocks):
+            blocks[self._cursor] = _RE.sub(
+                self._maybe_process, blocks[self._cursor]
+            )
+            self._cursor += 1
 
         # Return text unmodified, as we only need to modify the stashed raw HTML
         # blocks, which will later be reinstated by the raw HTML postprocessor