feat: Allow extensions to add templates

pawamoy · web-flow · commit cf0af059eb89 · 2023-05-24T12:04:41.000+02:00
An extension here is simply a Python package that defines an entry-point for a specific handler. For example, an extension can add templates to the Python handler thanks to this entry-point: ```toml [project.entry-points."mkdocstrings.python.templates"] extension-name = "extension_package:get_templates_path" ``` This entry-point assumes that the extension provides a `get_templates_path` function directly under the `extension_package` package. This function doesn't accept any argument and returns the path to a directory containing templates. The directory must contain one subfolder for each supported theme, for example: ``` templates/ material/ readthedocs/ mkdocs/ ``` mkdocstrings will add the folders corresponding to the user-selected theme, and to the handler's defined fallback theme, as usual, to the Jinja loader. The names of the extension templates must not overlap with the handler's original templates. The extension is then responsible, in collaboration with its target handler, for mutating the collected data in order to instruct the handler to use one of the extension template when rendering particular objects. For example, the Python handler will look for a `template` attribute on objects, and use it to render the object. This `template` attribute will be set by Griffe extensions (Griffe is the tool used by the Python handler to collect data). PR #569: #569
diff --git a/docs/usage/handlers.md b/docs/usage/handlers.md
@@ -192,7 +192,7 @@ is written for inspiration.
 ### Templates
 
 Your handler's implementation should normally be backed by templates, which go
-to the directory `mkdocstrings_handlers/custom_handler/templates/some_theme`.
+to the directory `mkdocstrings_handlers/custom_handler/templates/some_theme`
 (`custom_handler` here should be replaced with the actual name of your handler,
 and `some_theme` should be the name of an actual MkDocs theme that you support,
 e.g. `material`).
@@ -258,3 +258,79 @@ plugins:
       some_config_option: "b"
       other_config_option: 1
 ```
+
+## Handler extensions
+
+*mkdocstrings* provides a way for third-party packages
+to extend or alter the behavior of handlers.
+For example, an extension of the Python handler
+could add specific support for another Python library.
+
+NOTE: This feature is intended for developers.
+If you are a user and want to customize how objects are rendered,
+see [Theming / Customization](../theming/#customization).
+
+Such extensions can register additional template folders
+that will be used when rendering collected data.
+Extensions are responsible for synchronizing
+with the handler itself so that it uses the additional templates.
+
+An extension is a Python package
+that defines an entry-point for a specific handler:
+
+```toml title="pyproject.toml"
+[project.entry-points."mkdocstrings.python.templates"] # (1)!
+extension-name = "extension_package:get_templates_path" # (2)!
+```
+
+1. Replace `python` by the name of the handler you want to add templates to.
+1. Replace `extension-name` by any name you want,
+    and replace `extension_package:get_templates_path`
+    by the actual module path and function name in your package.
+
+This entry-point assumes that the extension provides
+a `get_templates_path` function directly under the `extension_package` package:
+
+```tree
+pyproject.toml
+extension_package/
+    __init__.py
+    templates/
+```
+
+```python title="extension_package/__init__.py"
+from pathlib import Path
+
+
+def get_templates_path() -> Path:
+    return Path(__file__).parent / "templates"
+```
+
+This function doesn't accept any argument
+and returns the path ([`pathlib.Path`][] or [`str`][])
+to a directory containing templates.
+The directory must contain one subfolder
+for each supported theme, even if empty
+(see "fallback theme" in [custom handlers templates](#templates_1)).
+For example:
+
+```tree
+pyproject.toml
+extension_package/
+    __init__.py
+    templates/
+        material/
+        readthedocs/
+        mkdocs/
+```
+
+*mkdocstrings* will add the folders corresponding to the user-selected theme,
+and to the handler's defined fallback theme, as usual.
+
+The names of the extension templates
+must not overlap with the handler's original templates.
+
+The extension is then responsible, in collaboration with its target handler,
+for mutating the collected data in order to instruct the handler
+to use one of the extension template when rendering particular objects.
+See each handler's docs to see if they support extensions, and how.
diff --git a/pyproject.toml b/pyproject.toml
@@ -35,6 +35,7 @@ dependencies = [
     "mkdocs>=1.2",
     "mkdocs-autorefs>=0.3.1",
     "pymdown-extensions>=6.3",
+    "importlib-metadata>=4.6; python_version < '3.10'",
     "typing-extensions>=4.1; python_version < '3.10'",
 ]
 
diff --git a/src/mkdocstrings/handlers/base.py b/src/mkdocstrings/handlers/base.py
@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import importlib
+import sys
 import warnings
 from contextlib import suppress
 from pathlib import Path
@@ -31,6 +32,12 @@
 from mkdocstrings.inventory import Inventory
 from mkdocstrings.loggers import get_template_logger
 
+# TODO: remove once support for Python 3.9 is dropped
+if sys.version_info < (3, 10):
+    from importlib_metadata import entry_points
+else:
+    from importlib.metadata import entry_points
+
 CollectorItem = Any
 
 
@@ -93,12 +100,23 @@ def __init__(self, handler: str, theme: str, custom_templates: str | None = None
         self._theme = theme
         self._custom_templates = custom_templates
 
+        # add selected theme templates
         themes_dir = self.get_templates_dir(handler)
         paths.append(themes_dir / theme)
 
+        # add extended theme templates
+        extended_templates_dirs = self.get_extended_templates_dirs(handler)
+        for templates_dir in extended_templates_dirs:
+            paths.append(templates_dir / theme)
+
+        # add fallback theme templates
         if self.fallback_theme and self.fallback_theme != theme:
             paths.append(themes_dir / self.fallback_theme)
 
+            # add fallback theme of extended templates
+            for templates_dir in extended_templates_dirs:
+                paths.append(templates_dir / self.fallback_theme)
+
         for path in paths:
             css_path = path / "style.css"
             if css_path.is_file():
@@ -179,6 +197,18 @@ def get_templates_dir(self, handler: str) -> Path:
 
         raise FileNotFoundError(f"Can't find 'templates' folder for handler '{handler}'")
 
+    def get_extended_templates_dirs(self, handler: str) -> list[Path]:
+        """Load template extensions for the given handler, return their templates directories.
+
+        Arguments:
+            handler: The name of the handler to get the extended templates directory of.
+
+        Returns:
+            The extensions templates directories.
+        """
+        discovered_extensions = entry_points(group=f"mkdocstrings.{handler}.templates")
+        return [extension.load()() for extension in discovered_extensions]
+
     def get_anchors(self, data: CollectorItem) -> tuple[str, ...] | set[str]:
         """Return the possible identifiers (HTML anchors) for a collected item.
 
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
@@ -2,10 +2,18 @@
 
 from __future__ import annotations
 
+from contextlib import suppress
+from pathlib import Path
+from typing import TYPE_CHECKING
+
 import pytest
+from jinja2.exceptions import TemplateNotFound
 from markdown import Markdown
 
-from mkdocstrings.handlers.base import Highlighter
+from mkdocstrings.handlers.base import BaseRenderer, Highlighter
+
+if TYPE_CHECKING:
+    from mkdocstrings.plugin import MkdocstringsPlugin
 
 
 @pytest.mark.parametrize("extension_name", ["codehilite", "pymdownx.highlight"])
@@ -43,3 +51,53 @@ def test_highlighter_basic(extension_name: str | None, inline: bool) -> None:
     actual = hl.highlight("import foo", language="python", inline=inline)
     assert "import" in actual
     assert "import foo" not in actual  # Highlighting has split it up.
+
+
+@pytest.fixture(name="extended_templates")
+def fixture_extended_templates(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path:  # noqa: D103
+    monkeypatch.setattr(BaseRenderer, "get_extended_templates_dirs", lambda self, handler: [tmp_path])
+    return tmp_path
+
+
+def test_extended_templates(extended_templates: Path, plugin: MkdocstringsPlugin) -> None:
+    """Test the extended templates functionality.
+
+    Parameters:
+        extended_templates: Temporary folder.
+        plugin: Instance of our plugin.
+    """
+    handler = plugin._handlers.get_handler("python")  # type: ignore[union-attr]
+
+    # assert mocked method added temp path to loader
+    search_paths = handler.env.loader.searchpath  # type: ignore[union-attr]
+    assert any(str(extended_templates) in path for path in search_paths)
+
+    # assert "new" template is not found
+    for path in search_paths:
+        # TODO: use missing_ok=True once support for Python 3.7 is dropped
+        with suppress(FileNotFoundError):
+            Path(path).joinpath("new.html").unlink()
+    with pytest.raises(expected_exception=TemplateNotFound):
+        handler.env.get_template("new.html")
+
+    # check precedence: base theme, base fallback theme, extended theme, extended fallback theme
+    # start with last one and go back up
+    handler.env.cache = None
+
+    extended_fallback_theme = extended_templates.joinpath(handler.fallback_theme)
+    extended_fallback_theme.mkdir()
+    extended_fallback_theme.joinpath("new.html").write_text("extended fallback new")
+    assert handler.env.get_template("new.html").render() == "extended fallback new"
+
+    extended_theme = extended_templates.joinpath("mkdocs")
+    extended_theme.mkdir()
+    extended_theme.joinpath("new.html").write_text("extended new")
+    assert handler.env.get_template("new.html").render() == "extended new"
+
+    base_fallback_theme = Path(search_paths[1])
+    base_fallback_theme.joinpath("new.html").write_text("base fallback new")
+    assert handler.env.get_template("new.html").render() == "base fallback new"
+
+    base_theme = Path(search_paths[0])
+    base_theme.joinpath("new.html").write_text("base new")
+    assert handler.env.get_template("new.html").render() == "base new"

Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,7 @@ dependencies = [`
`35`	`35`	`"mkdocs>=1.2",`
`36`	`36`	`"mkdocs-autorefs>=0.3.1",`
`37`	`37`	`"pymdown-extensions>=6.3",`
	`38`	`+ "importlib-metadata>=4.6; python_version < '3.10'",`
`38`	`39`	`"typing-extensions>=4.1; python_version < '3.10'",`
`39`	`40`	`]`
`40`	`41`