1from __future__ import annotations
  2
  3from typing import Any
  4
  5import jinja2
  6from jinja2 import nodes
  7from jinja2.ext import Extension
  8from jinja2.nodes import CallBlock, Node
  9from jinja2.parser import Parser
 10from jinja2.runtime import Context
 11from plain.runtime import settings
 12from plain.templates import register_template_extension
 13from plain.templates.jinja.extensions import InclusionTagExtension
 14
 15
 16@register_template_extension
 17class HTMXJSExtension(InclusionTagExtension):
 18    tags = {"htmx_js"}  # noqa: RUF012 — jinja2 types `tags` as an instance attribute; ClassVar here fails ty's LSP check
 19    template_name = "htmx/js.html"
 20
 21    def get_context(
 22        self, context: Context, *args: Any, **kwargs: Any
 23    ) -> dict[str, Any]:
 24        request = context.get("request")
 25        return {
 26            "DEBUG": settings.DEBUG,
 27            "extensions": kwargs.get("extensions", []),
 28            "csp_nonce": request.csp_nonce if request else None,
 29        }
 30
 31
 32class _FragmentFound(Exception):
 33    """Raised to short-circuit template rendering once the target fragment is found."""
 34
 35    def __init__(self, content: str) -> None:
 36        self.content = content
 37
 38
 39@register_template_extension
 40class HTMXFragmentExtension(Extension):
 41    tags = {"htmxfragment"}  # noqa: RUF012 — jinja2 types `tags` as an instance attribute; ClassVar here fails ty's LSP check
 42
 43    def parse(self, parser: Parser) -> Node:
 44        lineno = next(parser.stream).lineno
 45
 46        fragment_name = parser.parse_expression()
 47
 48        kwargs = []
 49
 50        while parser.stream.current.type != "block_end":
 51            if parser.stream.current.type == "name":
 52                key = parser.stream.current.value
 53                parser.stream.skip()
 54                parser.stream.expect("assign")
 55                value = parser.parse_expression()
 56                kwargs.append(nodes.Keyword(key, value))
 57
 58        body = parser.parse_statements(("name:endhtmxfragment",), drop_needle=True)
 59
 60        call = self.call_method(
 61            "_render_htmx_fragment",
 62            args=[fragment_name, nodes.ContextReference()],
 63            kwargs=kwargs,
 64        )
 65
 66        callblock = CallBlock(call, [], [], body)
 67        callblock.set_lineno(lineno)
 68
 69        return callblock
 70
 71    def _render_htmx_fragment(
 72        self, fragment_name: str, context: dict[str, Any], caller: Any, **kwargs: Any
 73    ) -> str:
 74        # Two-phase fragment targeting (see render_template_fragment):
 75        # Phase 1 skips non-target bodies, phase 2 renders them for nesting.
 76        # Once the target is found, "found" is set so child fragments render
 77        # normally with their wrapper divs.
 78        target_state = context.get("_htmx_target_fragment")
 79        if target_state is not None and not target_state["found"]:
 80            if str(fragment_name) == target_state["name"]:
 81                target_state["found"] = True
 82                content = caller()
 83                raise _FragmentFound(content)
 84            elif target_state["render_bodies"]:
 85                return caller()
 86            else:
 87                return ""
 88
 89        def attrs_to_str(attrs: dict[str, Any]) -> str:
 90            parts = []
 91            for k, v in attrs.items():
 92                if v == "":
 93                    parts.append(k)
 94                else:
 95                    parts.append(f'{k}="{v}"')
 96            return " ".join(parts)
 97
 98        render_lazy = kwargs.get("lazy", False)
 99        as_element = kwargs.get("as", "div")
100        attrs = {}
101        for k, v in kwargs.items():
102            if k in ("lazy", "as"):
103                continue
104            if k.startswith("hx_"):
105                attrs[k.replace("_", "-")] = v
106            else:
107                attrs[k] = v
108
109        if render_lazy:
110            attrs.setdefault("hx-trigger", "load from:body")
111            attrs.setdefault("hx-swap", "outerHTML")
112            attrs.setdefault("hx-target", "this")
113            attrs.setdefault("hx-indicator", "this")
114            attrs_str = attrs_to_str(attrs)
115            return f'<{as_element} plain-hx-fragment="{fragment_name}" hx-get {attrs_str}></{as_element}>'
116        else:
117            # Swap innerHTML so we can re-run hx calls inside the fragment automatically
118            attrs.setdefault("hx-swap", "innerHTML")
119            attrs.setdefault("hx-target", "this")
120            attrs.setdefault("hx-indicator", "this")
121            # Add an id that you can use to target the fragment from outside the fragment
122            attrs.setdefault("id", f"plain-hx-fragment-{fragment_name}")
123            attrs_str = attrs_to_str(attrs)
124            return f'<{as_element} plain-hx-fragment="{fragment_name}" {attrs_str}>{caller()}</{as_element}>'
125
126
127def render_template_fragment(
128    *, template: jinja2.Template, fragment_name: str, context: dict[str, Any]
129) -> str:
130    """Render only the named fragment from a template.
131
132    Two-phase approach:
133    1. Skip non-target fragment bodies (fast — handles top-level and loop fragments)
134    2. If not found, render bodies too (handles fragments nested inside other fragments)
135
136    Raises _FragmentFound to short-circuit as soon as the target is found.
137    """
138    for render_bodies in (False, True):
139        target_state = {
140            "name": fragment_name,
141            "found": False,
142            "render_bodies": render_bodies,
143        }
144        try:
145            template.render({**context, "_htmx_target_fragment": target_state})
146        except _FragmentFound as e:
147            return e.content
148
149    raise jinja2.TemplateNotFound(
150        f"Fragment '{fragment_name}' not found in template {template.name}"
151    )