1import os
  2import re
  3from typing import Any
  4
  5from jinja2 import nodes, pass_context
  6from jinja2.environment import Environment
  7from jinja2.ext import Extension
  8from jinja2.parser import Parser
  9from jinja2.runtime import Context
 10from plain.templates import register_template_extension
 11from plain.utils.safestring import SafeString, mark_safe
 12
 13
 14@pass_context
 15def Element(ctx: Context, _element_name: str, **kwargs: Any) -> SafeString:
 16    element_path_name = _element_name.replace(".", os.sep)
 17    template = ctx.environment.get_template(f"elements/{element_path_name}.html")
 18
 19    if "caller" in kwargs and "children" not in kwargs:
 20        # If we have a caller, we need to pass it as the children
 21        kwargs["children"] = kwargs["caller"]()
 22
 23    output = template.render(
 24        {
 25            # Note that this passes globals, but not things like loop variables
 26            # so for the most part you need to manually pass the kwargs you want
 27            **ctx.get_all(),
 28            **kwargs,
 29        }
 30    )
 31    return mark_safe(output)
 32
 33
 34@register_template_extension
 35class ElementsExtension(Extension):
 36    tags = {"use_elements"}  # noqa: RUF012 — jinja2 types `tags` as an instance attribute; ClassVar here fails ty's LSP check
 37
 38    def __init__(self, env: Environment):
 39        super().__init__(env)
 40        # Make the Element function available in connection with this extension
 41        env.globals["Element"] = Element  # ty: ignore[invalid-assignment]
 42
 43        self._CAP_TAG = r"(?:[a-z_]+\.)?[A-Z][A-Za-z0-9_]*"
 44
 45        # Jinja comments may mention capitalized tags as documentation —
 46        # mask them out before scanning so they're not treated as elements.
 47        self._COMMENT = re.compile(
 48            re.escape(env.comment_start_string)
 49            + r"[\s\S]*?"
 50            + re.escape(env.comment_end_string)
 51        )
 52
 53        self._SELF = re.compile(
 54            rf"<(?P<name>{self._CAP_TAG})(?P<attrs>(?:\s+(?:[^/>]|/(?!>))*?)?)/>"
 55        )
 56        self._CLOSED = re.compile(
 57            rf"<(?P<name>{self._CAP_TAG})(?P<attrs>(?:\s+[^>]*?)?)>"
 58            rf"(?P<body>[\s\S]*?)"
 59            rf"</\1>"
 60        )
 61
 62    def parse(self, parser: Parser) -> nodes.Output:
 63        # Consume {% use_elements %} and output nothing
 64        parser.stream.skip()
 65        return nodes.Output([])
 66
 67    def preprocess(
 68        self, source: str, name: str | None, filename: str | None = None
 69    ) -> str:
 70        if "{% use_elements %}" in source:
 71            # If we have a use_elements tag, we need to replace the template element tags
 72            # with the Element() calls
 73            source = self.replace_template_element_tags(source)
 74
 75        return source
 76
 77    def replace_template_element_tags(self, contents: str) -> str:
 78        if not contents:
 79            return contents
 80
 81        # Stash Jinja comments so capitalized tags inside them aren't
 82        # mistaken for elements, then restore them before returning.
 83        comments: list[str] = []
 84
 85        def stash_comment(m: re.Match) -> str:
 86            comments.append(m.group(0))
 87            return f"\x00element_comment_{len(comments) - 1}\x00"
 88
 89        contents = self._COMMENT.sub(stash_comment, contents)
 90
 91        def repl_self(m: re.Match) -> str:
 92            return self.convert_element(m.group("name"), m.group("attrs") or "", "")
 93
 94        def repl_closed(m: re.Match) -> str:
 95            body = m.group("body")
 96            if f"<{m.group('name')} " in body:
 97                raise ValueError(
 98                    f"Element {m.group('name')} cannot be nested in itself"
 99                )
100            return self.convert_element(m.group("name"), m.group("attrs") or "", body)
101
102        # keep stripping tags until we can’t find any more
103        prev = None
104        while prev != contents:
105            prev = contents
106            contents = self._SELF.sub(repl_self, contents)
107            contents = self._CLOSED.sub(repl_closed, contents)
108
109        if matches := re.search(rf"<{self._CAP_TAG}", contents):
110            raise ValueError(
111                f"Found unmatched capitalized tag in template: {matches.group(0)}"
112            )
113
114        if comments:
115            contents = re.sub(
116                r"\x00element_comment_(\d+)\x00",
117                lambda m: comments[int(m.group(1))],
118                contents,
119            )
120
121        return contents
122
123    def convert_element(self, element_name: str, s: str, children: str) -> str:
124        attrs: dict[str, str] = {}
125
126        # Quoted attrs
127        for k, v in re.findall(r'([a-zA-Z0-9_]+)="([^"]*)"', s):
128            attrs[k] = f'"{v}"'
129        for k, v in re.findall(r"([a-zA-Z0-9_]+)='([^']*)'", s):
130            attrs[k] = f"'{v}'"
131
132        # Bare attrs (assume they are strings)
133        for k, v in re.findall(r"([a-zA-Z0-9_]+)=([a-zA-Z0-9_\.]+)", s):
134            attrs[k] = f'"{v}"'
135
136        # Braced Python variables (remove the braces)
137        for k, raw in re.findall(r"([a-zA-Z0-9_]+)=({[^}]*})", s):
138            expr = raw[1:-1]
139            attrs[k] = expr
140
141        attrs_str = ", ".join(f"{k}={v}" for k, v in attrs.items())
142        if attrs_str:
143            attrs_str = ", " + attrs_str
144
145        call = f'{{% call Element("{element_name}"{attrs_str}) %}}{children}{{% endcall %}}'
146        return call.strip()