From cf0af059eb89240eba0437de417c124389e2f20e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timoth=C3=A9e=20Mazzucotelli?= <pawamoy@pm.me>
Date: Wed, 24 May 2023 12:04:41 +0200
Subject: [PATCH] feat: Allow extensions to add templates

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: https://github.com/mkdocstrings/mkdocstrings/pull/569
---
 docs/usage/handlers.md            | 78 ++++++++++++++++++++++++++++++-
 pyproject.toml                    |  1 +
 src/mkdocstrings/handlers/base.py | 30 ++++++++++++
 tests/test_handlers.py            | 60 +++++++++++++++++++++++-
 4 files changed, 167 insertions(+), 2 deletions(-)

diff --git a/docs/usage/handlers.md b/docs/usage/handlers.md
index 2083013c..e381090e 100644
--- 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
index 3c28fddf..a36f6d25 100644
--- 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
index 51efd775..b6ecd4fa 100644
--- 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
index a0b3be3e..777173e8 100644
--- 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"