matrix-ops-bot/ops_bot/util/template.py

# gitlab - A GitLab client and webhook receiver for maubot
# Copyright (C) 2021 Tulir Asokan
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
import os
import os.path
from pathlib import Path
from typing import Any, Callable, Dict, List, Tuple, Union

from jinja2 import BaseLoader
from jinja2 import Environment as JinjaEnvironment
from jinja2 import Template, TemplateNotFound

from ops_bot import get_project_root
from ops_bot.util import markdown


def sync_read_file(path: Union[Path, str]) -> str:
    with open(path) as file:
        return file.read()


def sync_list_files(directory: Union[Path, str]) -> list[str]:
    return os.listdir(directory)


class TemplateUtil:
    @staticmethod
    def bold_scope(label: str) -> str:
        try:
            scope, label = label.rsplit("::", 1)
            return f"{scope}::<strong>{label}</strong>"
        except ValueError:
            return label

    @staticmethod
    def pluralize(val: int, unit: str) -> str:
        if val == 1:
            return f"{val} {unit}"
        return f"{val} {unit}s"

    @classmethod
    def format_time(cls, seconds: Union[int, float], enable_days: bool = False) -> str:
        seconds = abs(seconds)
        frac_seconds = round(seconds - int(seconds), 1)
        minutes, seconds = divmod(int(seconds), 60)
        hours, minutes = divmod(minutes, 60)
        if enable_days:
            days, hours = divmod(hours, 24)
        else:
            days = 0
        parts = []
        if days > 0:
            parts.append(cls.pluralize(days, "day"))
        if hours > 0:
            parts.append(cls.pluralize(hours, "hour"))
        if minutes > 0:
            parts.append(cls.pluralize(minutes, "minute"))
        if seconds > 0 or len(parts) == 0:
            parts.append(cls.pluralize(int(seconds + frac_seconds), "second"))

        if len(parts) == 1:
            return parts[0]
        return ", ".join(parts[:-1]) + f" and {parts[-1]}"

    @staticmethod
    def join_human_list(
        data: List[str],
        *,
        joiner: str = ", ",
        final_joiner: str = " and ",
        mutate: Callable[[str], str] = lambda val: val,
    ) -> str:
        if not data:
            return ""
        elif len(data) == 1:
            return mutate(data[0])
        return (
            joiner.join(mutate(val) for val in data[:-1])
            + final_joiner
            + mutate(data[-1])
        )


class TemplateProxy:
    _env: JinjaEnvironment
    _args: Dict[str, Any]

    def __init__(self, env: JinjaEnvironment, args: Dict[str, Any]) -> None:
        self._env = env
        self._args = args

    def __getattr__(self, item: str) -> str:
        try:
            tpl = self._env.get_template(item)
        except TemplateNotFound:
            raise AttributeError(item)
        return tpl.render(**self._args)


class PluginTemplateLoader(BaseLoader):
    directory: Path
    macros: str

    def __init__(self, base: str, directory: str) -> None:
        template_root = os.environ.get("OPS_BOT_TEMPLATE_ROOT")
        if template_root is None:
            template_root = str(get_project_root() / ".." / "templates")
        template_root_path = Path(template_root)
        base_path = template_root_path / base
        self.directory = base_path / directory
        self.macros = sync_read_file(base_path / "macros.html")

    def get_source(
        self, environment: Any, template: str
    ) -> Tuple[str, str, Callable[[], bool]]:
        path = self.directory / f"{template}.html"
        try:
            tpl = sync_read_file(path)
        except KeyError:
            raise TemplateNotFound(template)
        return self.macros + tpl, template, lambda: True

    def list_templates(self) -> List[str]:
        return [
            os.path.splitext(os.path.basename(path))[0]
            for path in sync_list_files(self.directory)
            if path.endswith(".html")
        ]


class TemplateManager:
    _env: JinjaEnvironment
    _loader: PluginTemplateLoader

    def __init__(self, base: str, directory: str) -> None:
        self._loader = PluginTemplateLoader(base, directory)
        self._env = JinjaEnvironment(  # nosec B701
            loader=self._loader,
            lstrip_blocks=True,
            trim_blocks=True,
            autoescape=False,
            extensions=["jinja2.ext.do"],
        )
        self._env.filters["markdown"] = lambda message: markdown.render(
            message, allow_html=True
        )

    def __getitem__(self, item: str) -> Template:
        return self._env.get_template(item)

    def proxy(self, args: Dict[str, Any]) -> TemplateProxy:
        return TemplateProxy(self._env, args)