1from __future__ import annotations
2
3from typing import Any
4
5import jinja2
6from jinja2 import meta, nodes
7from jinja2.ext import Extension
8from jinja2.nodes import CallBlock, Node
9from jinja2.parser import Parser
10
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"}
19 template_name = "htmx/js.html"
20
21 def get_context(
22 self, context: dict[str, Any], *args: Any, **kwargs: Any
23 ) -> dict[str, Any]:
24 return {
25 "DEBUG": settings.DEBUG,
26 "extensions": kwargs.get("extensions", []),
27 }
28
29
30@register_template_extension
31class HTMXFragmentExtension(Extension):
32 tags = {"htmxfragment"}
33
34 def __init__(self, environment: jinja2.Environment):
35 super().__init__(environment)
36 environment.extend(htmx_fragment_nodes={})
37
38 def parse(self, parser: Parser) -> Node:
39 lineno = next(parser.stream).lineno
40
41 fragment_name = parser.parse_expression()
42
43 kwargs = []
44
45 while parser.stream.current.type != "block_end":
46 if parser.stream.current.type == "name":
47 key = parser.stream.current.value
48 parser.stream.skip()
49 parser.stream.expect("assign")
50 value = parser.parse_expression()
51 kwargs.append(nodes.Keyword(key, value))
52
53 body = parser.parse_statements(("name:endhtmxfragment",), drop_needle=True)
54
55 call = self.call_method(
56 "_render_htmx_fragment",
57 args=[fragment_name, nodes.ContextReference()],
58 kwargs=kwargs,
59 )
60
61 node = CallBlock(call, [], [], body).set_lineno(lineno)
62
63 # Store a reference to the node for later
64 self.environment.htmx_fragment_nodes.setdefault(parser.name, {})[ # type: ignore[attr-defined]
65 fragment_name.value # type: ignore[attr-defined]
66 ] = node
67
68 return node
69
70 def _render_htmx_fragment(
71 self, fragment_name: str, context: dict[str, Any], caller: Any, **kwargs: Any
72 ) -> str:
73 def attrs_to_str(attrs: dict[str, Any]) -> str:
74 parts = []
75 for k, v in attrs.items():
76 if v == "":
77 parts.append(k)
78 else:
79 parts.append(f'{k}="{v}"')
80 return " ".join(parts)
81
82 render_lazy = kwargs.get("lazy", False)
83 as_element = kwargs.get("as", "div")
84 attrs = {}
85 for k, v in kwargs.items():
86 if k.startswith("hx_"):
87 attrs[k.replace("_", "-")] = v
88 else:
89 attrs[k] = v
90
91 if render_lazy:
92 attrs.setdefault("hx-swap", "outerHTML")
93 attrs.setdefault("hx-target", "this")
94 attrs.setdefault("hx-indicator", "this")
95 attrs_str = attrs_to_str(attrs)
96 return f'<{as_element} plain-hx-fragment="{fragment_name}" hx-get hx-trigger="load from:body" {attrs_str}></{as_element}>'
97 else:
98 # Swap innerHTML so we can re-run hx calls inside the fragment automatically
99 # (render_template_fragment won't render this part of the node again, just the inner nodes)
100 attrs.setdefault("hx-swap", "innerHTML")
101 attrs.setdefault("hx-target", "this")
102 attrs.setdefault("hx-indicator", "this")
103 # Add an id that you can use to target the fragment from outside the fragment
104 attrs.setdefault("id", f"plain-hx-fragment-{fragment_name}")
105 attrs_str = attrs_to_str(attrs)
106 return f'<{as_element} plain-hx-fragment="{fragment_name}" {attrs_str}>{caller()}</{as_element}>'
107
108
109def render_template_fragment(
110 *, template: jinja2.Template, fragment_name: str, context: dict[str, Any]
111) -> str:
112 template = find_template_fragment(template, fragment_name)
113 return template.render(context)
114
115
116def find_template_fragment(
117 template: jinja2.Template, fragment_name: str
118) -> jinja2.Template:
119 # Look in this template for the fragment
120 callblock_node = template.environment.htmx_fragment_nodes.get( # type: ignore[attr-defined]
121 template.name, {}
122 ).get(fragment_name)
123
124 if not callblock_node:
125 # Look in other templates for this fragment
126 matching_callblock_nodes = []
127 for fragments in template.environment.htmx_fragment_nodes.values(): # type: ignore[attr-defined]
128 if fragment_name in fragments:
129 matching_callblock_nodes.append(fragments[fragment_name])
130
131 if len(matching_callblock_nodes) == 0:
132 # If we still haven't found anything, it's possible that we're
133 # in a different/new worker/process and haven't parsed the related templates yet
134 if template.environment.loader and template.name:
135 ast = template.environment.parse(
136 template.environment.loader.get_source(
137 template.environment, template.name
138 )[0]
139 )
140 for ref in meta.find_referenced_templates(ast):
141 if ref not in template.environment.htmx_fragment_nodes: # type: ignore[attr-defined]
142 # Trigger them to parse
143 template.environment.get_template(ref)
144
145 # Now look again
146 for fragments in template.environment.htmx_fragment_nodes.values(): # type: ignore[attr-defined]
147 if fragment_name in fragments:
148 matching_callblock_nodes.append(fragments[fragment_name])
149
150 if len(matching_callblock_nodes) == 1:
151 callblock_node = matching_callblock_nodes[0]
152 elif len(matching_callblock_nodes) > 1:
153 raise jinja2.TemplateNotFound(
154 f"Fragment {fragment_name} found in multiple templates. Use a more specific name."
155 )
156 else:
157 raise jinja2.TemplateNotFound(
158 f"Fragment {fragment_name} not found in any templates"
159 )
160
161 if not callblock_node:
162 raise jinja2.TemplateNotFound(
163 f"Fragment {fragment_name} not found in template {template.name}"
164 )
165
166 # Create a new template from the node
167 template_node = nodes.Template(callblock_node.body) # type: ignore[attr-defined]
168 return template.environment.from_string(template_node) # type: ignore[arg-type]