plugins.py

"""
Implements the plugin API for MkDocs.

"""
from __future__ import annotations

import logging
import sys
from collections import OrderedDict
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Generic,
    List,
    Optional,
    Tuple,
    Type,
    TypeVar,
    overload,
)

if sys.version_info >= (3, 10):
    from importlib.metadata import EntryPoint, entry_points
else:
    from importlib_metadata import EntryPoint, entry_points

if sys.version_info >= (3, 8):
    from typing import Literal
else:
    from typing_extensions import Literal

import jinja2.environment

from mkdocs import utils
from mkdocs.config.base import Config, ConfigErrors, ConfigWarnings, LegacyConfig, PlainConfigSchema
from mkdocs.livereload import LiveReloadServer
from mkdocs.structure.files import Files
from mkdocs.structure.nav import Navigation
from mkdocs.structure.pages import Page

if TYPE_CHECKING:
    from mkdocs.config.defaults import MkDocsConfig


log = logging.getLogger('mkdocs.plugins')


def get_plugins() -> Dict[str, EntryPoint]:
    """Return a dict of all installed Plugins as {name: EntryPoint}."""

    plugins = entry_points(group='mkdocs.plugins')

    # Allow third-party plugins to override core plugins
    pluginmap = {}
    for plugin in plugins:
        if plugin.name in pluginmap and plugin.value.startswith("mkdocs.contrib."):
            continue

        pluginmap[plugin.name] = plugin

    return pluginmap


SomeConfig = TypeVar('SomeConfig', bound=Config)


class BasePlugin(Generic[SomeConfig]):
    """
    Plugin base class.

    All plugins should subclass this class.
    """

    config_class: Type[SomeConfig] = LegacyConfig  # type: ignore[assignment]
    config_scheme: PlainConfigSchema = ()
    config: SomeConfig = {}  # type: ignore[assignment]

    def __class_getitem__(cls, config_class: Type[Config]):
        """Eliminates the need to write `config_class = FooConfig` when subclassing BasePlugin[FooConfig]"""
        name = f'{cls.__name__}[{config_class.__name__}]'
        return type(name, (cls,), dict(config_class=config_class))

    def __init_subclass__(cls):
        if not issubclass(cls.config_class, Config):
            raise TypeError(
                f"config_class {cls.config_class} must be a subclass of `mkdocs.config.base.Config`"
            )
        if cls.config_class is not LegacyConfig:
            cls.config_scheme = cls.config_class._schema  # For compatibility.

    def load_config(
        self, options: Dict[str, Any], config_file_path: Optional[str] = None
    ) -> Tuple[ConfigErrors, ConfigWarnings]:
        """Load config from a dict of options. Returns a tuple of (errors, warnings)."""

        if self.config_class is LegacyConfig:
            self.config = LegacyConfig(self.config_scheme, config_file_path=config_file_path)  # type: ignore
        else:
            self.config = self.config_class(config_file_path=config_file_path)

        self.config.load_dict(options)

        return self.config.validate()

    # One-time events

    def on_startup(self, *, command: Literal['build', 'gh-deploy', 'serve'], dirty: bool) -> None:
        """
        The `startup` event runs once at the very beginning of an `mkdocs` invocation.

        New in MkDocs 1.4.

        The presence of an `on_startup` method (even if empty) migrates the plugin to the new
        system where the plugin object is kept across builds within one `mkdocs serve`.

        Note that for initializing variables, the `__init__` method is still preferred.
        For initializing per-build variables (and whenever in doubt), use the `on_config` event.

        Parameters:
            command: the command that MkDocs was invoked with, e.g. "serve" for `mkdocs serve`.
            dirty: whether `--dirtyreload` or `--dirty` flags were passed.
        """

    def on_shutdown(self) -> None:
        """
        The `shutdown` event runs once at the very end of an `mkdocs` invocation, before exiting.

        This event is relevant only for support of `mkdocs serve`, otherwise within a
        single build it's undistinguishable from `on_post_build`.

        New in MkDocs 1.4.

        The presence of an `on_shutdown` method (even if empty) migrates the plugin to the new
        system where the plugin object is kept across builds within one `mkdocs serve`.

        Note the `on_post_build` method is still preferred for cleanups, when possible, as it has
        a much higher chance of actually triggering. `on_shutdown` is "best effort" because it
        relies on detecting a graceful shutdown of MkDocs.
        """

    def on_serve(
        self, server: LiveReloadServer, *, config: MkDocsConfig, builder: Callable
    ) -> Optional[LiveReloadServer]:
        """
        The `serve` event is only called when the `serve` command is used during
        development. It runs only once, after the first build finishes.
        It is passed the `Server` instance which can be modified before
        it is activated. For example, additional files or directories could be added
        to the list of "watched" files for auto-reloading.

        Parameters:
            server: `livereload.Server` instance
            config: global configuration object
            builder: a callable which gets passed to each call to `server.watch`

        Returns:
            `livereload.Server` instance
        """
        return server

    # Global events

    def on_config(self, config: MkDocsConfig) -> Optional[Config]:
        """
        The `config` event is the first event called on build and is run immediately
        after the user configuration is loaded and validated. Any alterations to the
        config should be made here.

        Parameters:
            config: global configuration object

        Returns:
            global configuration object
        """
        return config

    def on_pre_build(self, *, config: MkDocsConfig) -> None:
        """
        The `pre_build` event does not alter any variables. Use this event to call
        pre-build scripts.

        Parameters:
            config: global configuration object
        """

    def on_files(self, files: Files, *, config: MkDocsConfig) -> Optional[Files]:
        """
        The `files` event is called after the files collection is populated from the
        `docs_dir`. Use this event to add, remove, or alter files in the
        collection. Note that Page objects have not yet been associated with the
        file objects in the collection. Use [Page Events](plugins.md#page-events) to manipulate page
        specific data.

        Parameters:
            files: global files collection
            config: global configuration object

        Returns:
            global files collection
        """
        return files

    def on_nav(
        self, nav: Navigation, *, config: MkDocsConfig, files: Files
    ) -> Optional[Navigation]:
        """
        The `nav` event is called after the site navigation is created and can
        be used to alter the site navigation.

        Parameters:
            nav: global navigation object
            config: global configuration object
            files: global files collection

        Returns:
            global navigation object
        """
        return nav

    def on_env(
        self, env: jinja2.Environment, *, config: MkDocsConfig, files: Files
    ) -> Optional[jinja2.Environment]:
        """
        The `env` event is called after the Jinja template environment is created
        and can be used to alter the
        [Jinja environment](https://jinja.palletsprojects.com/en/latest/api/#jinja2.Environment).

        Parameters:
            env: global Jinja environment
            config: global configuration object
            files: global files collection

        Returns:
            global Jinja Environment
        """
        return env

    def on_post_build(self, *, config: MkDocsConfig) -> None:
        """
        The `post_build` event does not alter any variables. Use this event to call
        post-build scripts.

        Parameters:
            config: global configuration object
        """

    def on_build_error(self, error: Exception) -> None:
        """
        The `build_error` event is called after an exception of any kind
        is caught by MkDocs during the build process.
        Use this event to clean things up before MkDocs terminates. Note that any other
        events which were scheduled to run after the error will have been skipped. See
        [Handling Errors](plugins.md#handling-errors) for more details.

        Parameters:
            error: exception raised
        """

    # Template events

    def on_pre_template(
        self, template: jinja2.Template, *, template_name: str, config: MkDocsConfig
    ) -> Optional[jinja2.Template]:
        """
        The `pre_template` event is called immediately after the subject template is
        loaded and can be used to alter the template.

        Parameters:
            template: a Jinja2 [Template](https://jinja.palletsprojects.com/en/latest/api/#jinja2.Template) object
            template_name: string filename of template
            config: global configuration object

        Returns:
            a Jinja2 [Template](https://jinja.palletsprojects.com/en/latest/api/#jinja2.Template) object
        """
        return template

    def on_template_context(
        self, context: Dict[str, Any], *, template_name: str, config: MkDocsConfig
    ) -> Optional[Dict[str, Any]]:
        """
        The `template_context` event is called immediately after the context is created
        for the subject template and can be used to alter the context for that specific
        template only.

        Parameters:
            context: dict of template context variables
            template_name: string filename of template
            config: global configuration object

        Returns:
            dict of template context variables
        """
        return context

    def on_post_template(
        self, output_content: str, *, template_name: str, config: MkDocsConfig
    ) -> Optional[str]:
        """
        The `post_template` event is called after the template is rendered, but before
        it is written to disc and can be used to alter the output of the template.
        If an empty string is returned, the template is skipped and nothing is is
        written to disc.

        Parameters:
            output_content: output of rendered template as string
            template_name: string filename of template
            config: global configuration object

        Returns:
            output of rendered template as string
        """
        return output_content

    # Page events

    def on_pre_page(self, page: Page, *, config: MkDocsConfig, files: Files) -> Optional[Page]:
        """
        The `pre_page` event is called before any actions are taken on the subject
        page and can be used to alter the `Page` instance.

        Parameters:
            page: `mkdocs.nav.Page` instance
            config: global configuration object
            files: global files collection

        Returns:
            `mkdocs.nav.Page` instance
        """
        return page

    def on_page_read_source(self, *, page: Page, config: MkDocsConfig) -> Optional[str]:
        """
        The `on_page_read_source` event can replace the default mechanism to read
        the contents of a page's source from the filesystem.

        Parameters:
            page: `mkdocs.nav.Page` instance
            config: global configuration object

        Returns:
            The raw source for a page as unicode string. If `None` is returned, the
                default loading from a file will be performed.
        """
        return None

    def on_page_markdown(
        self, markdown: str, *, page: Page, config: MkDocsConfig, files: Files
    ) -> Optional[str]:
        """
        The `page_markdown` event is called after the page's markdown is loaded
        from file and can be used to alter the Markdown source text. The meta-
        data has been stripped off and is available as `page.meta` at this point.

        Parameters:
            markdown: Markdown source text of page as string
            page: `mkdocs.nav.Page` instance
            config: global configuration object
            files: global files collection

        Returns:
            Markdown source text of page as string
        """
        return markdown

    def on_page_content(
        self, html: str, *, page: Page, config: MkDocsConfig, files: Files
    ) -> Optional[str]:
        """
        The `page_content` event is called after the Markdown text is rendered to
        HTML (but before being passed to a template) and can be used to alter the
        HTML body of the page.

        Parameters:
            html: HTML rendered from Markdown source as string
            page: `mkdocs.nav.Page` instance
            config: global configuration object
            files: global files collection

        Returns:
            HTML rendered from Markdown source as string
        """
        return html

    def on_page_context(
        self, context: Dict[str, Any], *, page: Page, config: MkDocsConfig, nav: Navigation
    ) -> Optional[Dict[str, Any]]:
        """
        The `page_context` event is called after the context for a page is created
        and can be used to alter the context for that specific page only.

        Parameters:
            context: dict of template context variables
            page: `mkdocs.nav.Page` instance
            config: global configuration object
            nav: global navigation object

        Returns:
            dict of template context variables
        """
        return context

    def on_post_page(self, output: str, *, page: Page, config: MkDocsConfig) -> Optional[str]:
        """
        The `post_page` event is called after the template is rendered, but
        before it is written to disc and can be used to alter the output of the
        page. If an empty string is returned, the page is skipped and nothing is
        written to disc.

        Parameters:
            output: output of rendered template as string
            page: `mkdocs.nav.Page` instance
            config: global configuration object

        Returns:
            output of rendered template as string
        """
        return output


EVENTS = tuple(k[3:] for k in BasePlugin.__dict__ if k.startswith("on_"))

# The above definitions were just for docs and type checking, we don't actually want them.
for k in EVENTS:
    delattr(BasePlugin, 'on_' + k)


T = TypeVar('T')


def event_priority(priority: float) -> Callable[[T], T]:
    """A decorator to set an event priority for an event handler method.

    Recommended priority values:
    `100` "first", `50` "early", `0` "default", `-50` "late", `-100` "last".
    As different plugins discover more precise relations to each other, the values should be further tweaked.

    ```python
    @plugins.event_priority(-100)  # Wishing to run this after all other plugins' `on_files` events.
    def on_files(self, files, config, **kwargs):
        ...
    ```

    New in MkDocs 1.4.
    Recommended shim for backwards compatibility:

    ```python
    try:
        from mkdocs.plugins import event_priority
    except ImportError:
        event_priority = lambda priority: lambda f: f  # No-op fallback
    ```
    """

    def decorator(event_method):
        event_method.mkdocs_priority = priority
        return event_method

    return decorator


class PluginCollection(OrderedDict):
    """
    A collection of plugins.

    In addition to being a dict of Plugin instances, each event method is registered
    upon being added. All registered methods for a given event can then be run in order
    by calling `run_event`.
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.events: Dict[str, List[Callable]] = {k: [] for k in EVENTS}

    def _register_event(self, event_name: str, method: Callable) -> None:
        """Register a method for an event."""
        utils.insort(
            self.events[event_name], method, key=lambda m: -getattr(m, 'mkdocs_priority', 0)
        )

    def __setitem__(self, key: str, value: BasePlugin, **kwargs) -> None:
        super().__setitem__(key, value, **kwargs)
        # Register all of the event methods defined for this Plugin.
        for event_name in (x for x in dir(value) if x.startswith('on_')):
            method = getattr(value, event_name, None)
            if callable(method):
                self._register_event(event_name[3:], method)

    @overload
    def run_event(self, name: str, item: None = None, **kwargs) -> Any:
        ...

    @overload
    def run_event(self, name: str, item: T, **kwargs) -> T:
        ...

    def run_event(self, name: str, item=None, **kwargs) -> Optional[T]:
        """
        Run all registered methods of an event.

        `item` is the object to be modified or replaced and returned by the event method.
        If it isn't given the event method creates a new object to be returned.
        All other keywords are variables for context, but would not generally
        be modified by the event method.
        """
        pass_item = item is not None
        events = self.events[name]
        if events:
            log.debug(f'Running {len(events)} `{name}` events')
        for method in events:
            if pass_item:
                result = method(item, **kwargs)
            else:
                result = method(**kwargs)
            # keep item if method returned `None`
            if result is not None:
                item = result
        return item