Coverage for jinja2_async_environment / compiler_modules / codegen.py: 81%
588 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-11-26 21:26 -0800
1"""Async code generator for template compilation."""
3import typing as t
5from jinja2 import nodes
6from jinja2.compiler import (
7 CodeGenerator,
8 Frame,
9 find_undeclared,
10)
12from .cache import CompilationCache
13from .dependencies import DependencyResolver
14from .frame import AsyncFrame
15from .loops import LoopCodeGenerator
16from .patterns import CompiledPatterns
18if t.TYPE_CHECKING:
19 pass
21# Global compilation cache instance
22_compilation_cache = CompilationCache()
25class AsyncCodeGenerator(CodeGenerator):
26 """Async-aware code generator extending Jinja2's CodeGenerator."""
28 environment: t.Any
29 name: str
30 filename: str
31 stream: t.Any
32 extends_so_far: int
33 has_known_extends: bool
34 root_frame_class: type[AsyncFrame] = AsyncFrame
35 eval_ctx: t.Any = None
36 is_async: bool = True
37 last_identifier: int = 0
38 identifiers: dict[str, t.Any] = {}
39 import_aliases: dict[str, t.Any] = {}
40 blocks: dict[str, t.Any] = {}
41 extends_buffer: t.Any = None
42 required_blocks: set[str] = set()
43 has_super: bool = False
44 macro_frames: list[AsyncFrame] = []
46 # Fast lookup cache for common variable names
47 _COMMON_VARS = frozenset(
48 [
49 "context",
50 "environment",
51 "eval_ctx",
52 "undefined",
53 "item",
54 "loop",
55 "block",
56 "value",
57 "name",
58 "key",
59 ]
60 )
62 def __init__(
63 self, environment: t.Any, name: str, filename: str, defer_init: bool = False
64 ) -> None:
65 super().__init__(
66 environment, name, filename, stream=None, defer_init=defer_init
67 )
68 self.extends_so_far = 0
69 self.has_known_extends = False
70 self.has_super = False
71 self.last_identifier = 0
72 self.identifiers = {}
73 self.import_aliases = {}
74 self.blocks = {}
75 self.extends_buffer = None
76 self.required_blocks = set()
77 self.is_async = True
78 self.macro_frames = []
80 # Initialize assignment tracking stack
81 self._assign_stack: list[set[str]] = []
83 # Initialize utility classes for better code organization
84 self._dependency_resolver = DependencyResolver(self)
85 self._loop_generator = LoopCodeGenerator(self)
87 from jinja2.nodes import EvalContext
89 if self.eval_ctx is None:
90 self.eval_ctx = EvalContext(self.environment, self.name)
92 def choose_async(self, async_fmt: str = "async ", sync_fmt: str = "") -> str: # type: ignore[override]
93 return async_fmt if self.environment.enable_async else sync_fmt
95 def simple_write(self, value: str, frame: Frame) -> None: # type: ignore[override]
96 self.writeline(f"yield {value}")
98 def func_code_generator(self, frame: Frame) -> str:
99 async_frame = t.cast(AsyncFrame, frame)
100 return "async def" if async_frame.is_async else "def"
102 def func(self, name: str) -> str:
103 """Generate a function declaration for the given name.
105 Properly handles async functions by checking environment.is_async.
106 """
107 return f"{self.choose_async()}def {name}"
109 def enter_frame(self, frame: Frame) -> None:
110 """Enter a new frame context and load variables from context.
112 This method generates code to resolve template variables from the
113 context and assign them to local frame variables (e.g., l_0_name).
114 """
115 from jinja2.compiler import (
116 VAR_LOAD_ALIAS, # type: ignore[attr-defined]
117 VAR_LOAD_PARAMETER, # type: ignore[attr-defined]
118 VAR_LOAD_RESOLVE, # type: ignore[attr-defined]
119 VAR_LOAD_UNDEFINED, # type: ignore[attr-defined]
120 )
122 undefs = []
123 for target, (action, param) in frame.symbols.loads.items():
124 if action == VAR_LOAD_PARAMETER:
125 pass
126 elif action == VAR_LOAD_RESOLVE:
127 self.writeline(f"{target} = {self.get_resolve_func()}({param!r})")
128 elif action == VAR_LOAD_ALIAS:
129 self.writeline(f"{target} = {param}")
130 elif action == VAR_LOAD_UNDEFINED:
131 undefs.append(target)
132 else:
133 raise NotImplementedError("unknown load instruction")
134 if undefs:
135 self.writeline(f"{' = '.join(undefs)} = missing")
137 def leave_frame(self, frame: Frame, with_python_scope: bool = False) -> None:
138 """Leave a frame context."""
139 pass
141 def return_buffer_contents(
142 self,
143 frame: Frame,
144 force_unescaped: bool = False, # noqa: ARG002
145 ) -> None:
146 _ = force_unescaped
147 if frame.buffer is not None:
148 self.writeline(f"return ''.join({frame.buffer})")
150 def visit_Name(self, node: nodes.Name, frame: Frame) -> None:
151 frame = t.cast(AsyncFrame, frame)
152 self._handle_assignment_tracking(node, frame)
153 if self._handle_special_names(node):
154 return
155 self._handle_symbol_name(node, frame)
157 def _handle_assignment_tracking(self, node: nodes.Name, frame: AsyncFrame) -> None:
158 if node.ctx == "store":
159 frame.symbols.store(node.name)
160 if frame.toplevel or frame.loop_frame or frame.block_frame:
161 if hasattr(self, "_assign_stack") and self._assign_stack:
162 self._assign_stack[-1].add(node.name)
164 def _handle_special_names(self, node: nodes.Name) -> bool:
165 if node.name in ("blocks", "debug_info"):
166 self.write(node.name)
167 return True
168 return False
170 def _handle_symbol_name(self, node: nodes.Name, frame: AsyncFrame) -> None:
171 # Fast path for common variables
172 if node.name in self._COMMON_VARS and node.ctx == "load":
173 try:
174 ref = frame.symbols.ref(node.name)
175 # Add undefined check for ALL variables, including common ones
176 # This ensures proper undefined variable handling
177 self.write(
178 f"(undefined(name={node.name!r}) if {ref} is missing else {ref})"
179 )
180 return
181 except AssertionError:
182 self.write(f"context.get({node.name!r})")
183 return
185 # Standard path for other variables
186 try:
187 ref = frame.symbols.ref(node.name)
188 if node.ctx == "load":
189 # ALWAYS use undefined check for variables loaded from context
190 # This ensures missing variables render as Undefined() instead of "missing"
191 self.write(
192 f"(undefined(name={node.name!r}) if {ref} is missing else {ref})"
193 )
194 else:
195 self.write(ref)
196 except AssertionError:
197 if node.ctx == "load":
198 self.write(f"context.get({node.name!r})")
199 else:
200 self.write(f"context.vars[{node.name!r}]")
202 def _should_use_undefined_check(self, ref: str, frame: AsyncFrame) -> bool:
203 # Try to get VAR_LOAD_PARAMETER using getattr to avoid import issues
204 try:
205 import jinja2.compiler
207 VAR_LOAD_PARAMETER = getattr(jinja2.compiler, "VAR_LOAD_PARAMETER", None)
208 except (ImportError, AttributeError):
209 # If we can't access it, we can't use it
210 return False
212 if VAR_LOAD_PARAMETER is None:
213 return False
215 load = frame.symbols.find_load(ref)
216 return not (
217 load is not None
218 and load[0] == VAR_LOAD_PARAMETER
219 and hasattr(self, "parameter_is_undeclared")
220 and not self.parameter_is_undeclared(ref)
221 )
223 def pull_dependencies(self, nodes: t.Iterable[nodes.Node]) -> None:
224 """Find all filter and test names used in the template and assign them to variables."""
225 from jinja2.compiler import DependencyFinderVisitor
227 visitor = DependencyFinderVisitor()
228 for node in nodes:
229 visitor.visit(node)
231 # Set up filter dependencies using utility class
232 for name in sorted(visitor.filters):
233 self._dependency_resolver.setup_filter_dependency(name)
235 # Set up test dependencies using utility class
236 for name in sorted(visitor.tests):
237 self._dependency_resolver.setup_test_dependency(name)
239 def generate(self, node: nodes.Template) -> str:
240 """Generate template code following base Jinja2 architecture.
242 Architecture (matching base Jinja2):
243 1. Pre-discover all blocks and store in self.blocks
244 2. Generate imports and module-level setup
245 3. Generate root function with block CALLS
246 4. After root exits, generate all block DEFINITIONS
247 5. Generate module-level blocks dict
249 This separation of block calling vs. definition is CRITICAL for
250 template inheritance to work correctly.
251 """
252 from jinja2.nodes import EvalContext
254 if self.eval_ctx is None:
255 self.eval_ctx = EvalContext(self.environment, self.name)
257 # PHASE 1: Pre-discover all blocks (like base Jinja2)
258 # This must happen BEFORE any code generation so blocks are known
259 for block in node.find_all(nodes.Block):
260 if block.name in self.blocks:
261 self.fail(f"block {block.name!r} defined twice", block.lineno)
262 self.blocks[block.name] = block
264 # Check if template has extends (for proper frame configuration)
265 have_extends = node.find(nodes.Extends) is not None
267 # PHASE 2: Module-level setup
268 # Use optimized cached imports for better performance
269 for import_line in CompiledPatterns.get_optimized_imports().split("\n"):
270 self.writeline(import_line)
272 self.writeline(f"name = {self.name!r}")
274 # Helper functions
275 self.writeline("def undefined(name=None, **_):")
276 self.indent()
277 self.writeline("return Undefined(name=name)")
278 self.outdent()
280 self.writeline("async def auto_await(value):")
281 self.indent()
282 self.writeline("if hasattr(value, '__await__'):")
283 self.indent()
284 self.writeline("return await value")
285 self.outdent()
286 self.writeline("return value")
287 self.outdent()
289 self.writeline("filters = DEFAULT_FILTERS.copy()")
290 self.writeline("filters['escape'] = escape")
292 # PHASE 3: Generate root render function
293 self.writeline("async def root(context):")
294 self.indent()
296 # CRITICAL: Write commons sets up resolve, undefined, concat, etc.
297 # AND includes "if 0: yield None" to make function an async generator
298 self.writeline("resolve = context.resolve_or_missing")
299 self.writeline("undefined = environment.undefined")
300 self.writeline("concat = environment.concat")
301 self.writeline("cond_expr_undefined = Undefined")
302 self.writeline("if 0: yield None") # Makes function an async generator
304 # Create frame for root processing
305 frame = self.root_frame_class(eval_ctx=self.eval_ctx)
307 # Check for 'self' usage BEFORE analyzing
308 if "self" in find_undeclared(node.body, ("self",)):
309 ref = frame.symbols.declare_parameter("self")
310 self.writeline(f"{ref} = TemplateReference(context)")
312 # Analyze node to discover variables
313 frame.symbols.analyze_node(node)
315 # Set frame flags AFTER analysis
316 frame.toplevel = frame.rootlevel = True
317 frame.require_output_check = have_extends and not self.has_known_extends
319 # Initialize parent_template for extends
320 if have_extends:
321 self.writeline("parent_template = None")
323 # Enter frame and process template body
324 # visit_Block will be called during blockvisit, but now it only CALLS blocks
325 self.enter_frame(frame)
326 self.pull_dependencies(node.body)
327 self.blockvisit(node.body, frame)
328 self.leave_frame(frame, with_python_scope=True)
329 self.outdent() # Exit root function
331 # PHASE 4: Parent template iteration (if extends present)
332 if have_extends:
333 if not self.has_known_extends:
334 self.indent()
335 self.writeline("if parent_template is not None:")
336 self.indent()
337 # Generate async parent template iteration
338 self.writeline("agen = parent_template.root_render_func(context)")
339 self.writeline("try:")
340 self.indent()
341 self.writeline("async for event in agen:")
342 self.indent()
343 self.writeline("yield event")
344 self.outdent()
345 self.outdent()
346 self.writeline("finally: await agen.aclose()")
347 self.outdent(1 + (not self.has_known_extends))
349 # PHASE 5: Generate all block function DEFINITIONS
350 # This happens AFTER root function exits
351 self._generate_block_functions()
353 # PHASE 6: Generate module-level blocks dict
354 blocks_kv_str = ", ".join(f"{name!r}: block_{name}" for name in self.blocks)
355 self.writeline(f"blocks = {{{blocks_kv_str}}}", extra=1)
356 self.writeline("debug_info = None")
358 # Apply pattern-based optimizations to generated code
359 generated_code = self.stream.getvalue()
360 return CompiledPatterns.optimize_generated_code(generated_code)
362 def _generate_block_functions(self) -> None:
363 """Generate all block function definitions.
365 This is called AFTER the root function exits, following base Jinja2's
366 architecture. Block functions are defined separately from where they're
367 called to enable proper template inheritance.
369 For each block in self.blocks:
370 - Creates async def block_NAME(context) function
371 - Adds write_commons() for generator setup
372 - Processes block body with proper frame
373 - Handles 'self' and 'super' references
374 """
375 for name, block_node in self.blocks.items():
376 # Start block function definition
377 self.writeline(
378 f"async def block_{name}(context):",
379 block_node,
380 1,
381 )
382 self.indent()
384 # Write commons (includes "if 0: yield None" for generator)
385 self.writeline("resolve = context.resolve_or_missing")
386 self.writeline("undefined = environment.undefined")
387 self.writeline("concat = environment.concat")
388 self.writeline("cond_expr_undefined = Undefined")
389 self.writeline("if 0: yield None") # Makes function an async generator
391 # Create block frame (NOT a child of toplevel frame)
392 block_frame = self.root_frame_class(eval_ctx=self.eval_ctx)
393 block_frame.block_frame = True
395 # Check for 'self' and 'super' usage
396 undeclared = find_undeclared(block_node.body, ("self", "super"))
397 if "self" in undeclared:
398 ref = block_frame.symbols.declare_parameter("self")
399 self.writeline(f"{ref} = TemplateReference(context)")
400 if "super" in undeclared:
401 ref = block_frame.symbols.declare_parameter("super")
402 self.writeline(f"{ref} = context.super({name!r}, block_{name})")
404 # Analyze block body
405 block_frame.symbols.analyze_node(block_node)
406 block_frame.block = name
408 # Block-level variables dict
409 self.writeline("_block_vars = {}")
411 # Enter frame, process block body, leave frame
412 self.enter_frame(block_frame)
413 self.pull_dependencies(block_node.body)
414 self.blockvisit(block_node.body, block_frame)
415 self.leave_frame(block_frame, with_python_scope=True)
417 self.outdent() # Exit block function
419 def visit_Block(self, node: nodes.Block, frame: Frame) -> None:
420 """Call a block and register it for the template.
422 This method ONLY handles CALLING blocks from context.blocks.
423 Block DEFINITIONS are generated separately in _generate_block_functions().
425 This separation is CRITICAL for template inheritance to work:
426 - Child templates can override parent blocks via context.blocks
427 - Blocks are discovered and defined after root function exits
428 - During root execution, blocks are called from context.blocks[name][0]
430 Architecture matches base Jinja2's CodeGenerator.visit_Block.
431 """
432 block_name = node.name
434 # Handle inheritance conditional
435 level = 0
436 if frame.toplevel:
437 # If we know we're a child template, no need to check
438 if self.has_known_extends:
439 return
440 # If we've seen extends before, add conditional
441 if self.extends_so_far > 0:
442 self.writeline("if parent_template is None:")
443 self.indent()
444 level += 1
446 # Determine context reference (scoped blocks use derived context)
447 if node.scoped:
448 context = self.derive_context(frame)
449 else:
450 context = self.get_context_ref()
452 # Check if block is required
453 if node.required:
454 self.writeline(f"if len(context.blocks[{block_name!r}]) <= 1:", node)
455 self.indent()
456 self.writeline(
457 f'raise TemplateRuntimeError("Required block {block_name!r} not found")',
458 node,
459 )
460 self.outdent()
462 # CRITICAL: Call block from context.blocks (enables inheritance)
463 # context.blocks[name] is a list, [0] gets the first (child's override or base)
464 self.writeline(f"gen = context.blocks[{block_name!r}][0]({context})")
465 self.writeline("try:")
466 self.indent()
467 self.writeline("async for event in gen:")
468 self.indent()
469 self.simple_write("event", frame)
470 self.outdent()
471 self.outdent()
472 self.writeline("finally: await gen.aclose()")
474 # Close inheritance conditional if needed
475 self.outdent(level)
477 def visit_Extends(self, node: nodes.Extends, frame: Frame) -> None:
478 """Visit an extends node with proper async handling."""
479 # Frame is already cast to AsyncFrame through method signature compatibility
481 # If output check is not required, raise CompilerExit immediately
482 if not frame.require_output_check:
483 from jinja2.compiler import CompilerExit
485 raise CompilerExit()
487 # Check if we're in a top-level scope
488 if not frame.toplevel:
489 self.fail("cannot use extend from a non top-level scope", node.lineno)
491 # Handle multiple extends
492 if self.extends_so_far > 0:
493 if not self.has_known_extends:
494 self.writeline("if parent_template is not None:")
495 self.indent()
496 self.writeline('raise TemplateRuntimeError("extended multiple times")')
497 if self.has_known_extends:
498 from jinja2.compiler import CompilerExit
500 raise CompilerExit()
501 else:
502 self.outdent()
504 # Generate async template loading code
505 self.writeline("parent_template = await environment.get_template_async(", node)
506 self.visit(node.template, frame)
507 self.write(f", {self.name!r})")
508 self.writeline("for name, parent_block in parent_template.blocks.items():")
509 self.indent()
510 self.writeline("context.blocks.setdefault(name, []).append(parent_block)")
511 self.outdent()
513 # Update inheritance tracking
514 if frame.rootlevel:
515 self.has_known_extends = True
516 self.extends_so_far += 1
518 def visit_Include(self, node: nodes.Include, frame: Frame) -> None:
519 """Visit an include node with proper async handling."""
520 # Frame is already cast to AsyncFrame through method signature compatibility
522 # Handle ignore_missing flag
523 if node.ignore_missing:
524 self.writeline("try:")
525 self.indent()
527 # Generate async template loading code
528 self.writeline("template = await environment.get_template_async(", node)
529 self.visit(node.template, frame)
530 self.write(f", {self.name!r})")
532 # Close try block for ignore_missing
533 if node.ignore_missing:
534 self.outdent()
535 self.writeline("except TemplateNotFound:")
536 self.indent()
537 self.writeline("pass")
538 self.outdent()
539 self.writeline("else:")
540 self.indent()
542 # Generate rendering code based on context flag
543 if node.with_context:
544 # With context - include local variables
545 local_context = self.dump_local_context(frame)
546 self.writeline(
547 f"async for event in template.root_render_func(template.new_context(context.get_all(), True, {local_context})):"
548 )
549 else:
550 # Without context - use default module
551 self.writeline(
552 "async for event in (await template._get_default_module_async())._body_stream:"
553 )
555 # Generate event output
556 self.indent()
557 self.simple_write("event", frame)
558 self.outdent()
560 # Close else block for ignore_missing
561 if node.ignore_missing:
562 self.outdent()
564 def visit_AsyncFor(self, node: nodes.For, frame: Frame) -> None:
565 """Visit an async for loop node with proper async handling."""
566 # Frame is already cast to AsyncFrame through method signature compatibility
568 # Handle recursive loops (not supported)
569 if node.recursive:
570 raise NotImplementedError("Recursive loops not supported")
572 # Get target variable name
573 target = node.target
574 if isinstance(target, nodes.Name):
575 item = target.name
576 else:
577 item = "item"
578 frame.symbols.store(item)
580 # Initialize target variable
581 self.writeline(f"{item} = None")
583 # Handle loop filter
584 loop_filter = None
585 if hasattr(node, "test") and node.test:
586 loop_filter = self.temporary_identifier()
587 self.writeline(f"{loop_filter} = ", node.test)
588 self.visit(node.test, frame)
590 # Initialize loop counter
591 loop_var = self.temporary_identifier()
592 self.writeline(f"{loop_var} = -1", node)
594 # Generate async for loop
595 self.writeline(f"async for {item} in ", node.iter)
596 self.visit(node.iter, frame)
597 self.write(":")
598 self.indent()
600 # Increment loop counter
601 self.writeline(f"{loop_var} += 1")
603 # Handle loop filter condition
604 if hasattr(node, "test") and node.test and loop_filter:
605 self.writeline(f"if {loop_filter}({item}):")
606 self.indent()
608 # Process loop body
609 if hasattr(node, "body"):
610 self.blockvisit(node.body, frame)
612 # Close filter condition
613 if hasattr(node, "test") and node.test and loop_filter:
614 self.outdent()
616 # Close main loop
617 self.outdent()
619 # Handle else clause
620 if hasattr(node, "else_") and node.else_:
621 self.writeline(f"if {loop_var} == -1:")
622 self.indent()
623 self.blockvisit(node.else_, frame)
624 self.outdent()
626 def visit_AsyncCall(self, node: nodes.Call, frame: Frame) -> None:
627 """Visit an async call node by adding await prefix."""
628 self.write("await ")
629 self.visit_Call(node, frame)
631 def visit_AsyncFilterBlock(self, node: nodes.FilterBlock, frame: Frame) -> None:
632 """Visit an async filter block node."""
633 # Frame is already cast to AsyncFrame through method signature compatibility
635 # Early return if no filter or body
636 if not hasattr(node, "filter"):
637 return
638 if not hasattr(node, "body"):
639 return
641 # Get filter node
642 filter_node = node.filter
644 # Create buffer for collecting content
645 buffer = self.temporary_identifier()
646 self.writeline(f"{buffer} = []")
648 # Create async frame for processing body
649 asyncframe = frame.copy()
650 asyncframe.buffer = buffer
651 asyncframe.toplevel = False
653 # Process the body
654 self.blockvisit(node.body, asyncframe)
656 # Generate await call for filter
657 self.writeline("await ", filter_node)
658 self.visit(filter_node, frame)
659 self.write(f"(''.join({buffer}))")
661 def visit_AsyncBlock(self, node: nodes.Block, frame: Frame) -> None:
662 """Visit an async block node."""
663 # Frame is already cast to AsyncFrame through method signature compatibility
665 # Early return if no name or body
666 if not hasattr(node, "name"):
667 return
668 if not hasattr(node, "body"):
669 return
671 # Get block name
672 block_name = node.name
674 # Initialize block storage
675 self.writeline(f"blocks[{block_name!r}] = []")
677 # Define async block function
678 block_func_name = f"block_{block_name}"
679 self.writeline(f"async def {block_func_name}(context):")
680 self.indent()
682 # Empty block content placeholder
683 self.writeline("yield ''")
685 # Process block body if it exists
686 if node.body:
687 self.blockvisit(node.body, frame)
689 # End function and register block
690 self.outdent()
691 self.writeline(f"blocks[{block_name!r}].append({block_func_name})")
693 def _import_common(
694 self, node: nodes.Import | nodes.FromImport, frame: Frame
695 ) -> None:
696 """Common import functionality with async template loading."""
697 # Cast frame to AsyncFrame for type safety
698 frame = t.cast(AsyncFrame, frame)
700 # Generate async template loading code
701 self.writeline("template = await environment.get_template_async(", node)
702 self.visit(node.template, frame)
703 self.write(f", {self.name!r})")
705 @classmethod
706 def compile_with_cache(
707 cls, environment: t.Any, source: str, name: str, filename: str
708 ) -> str:
709 """Compile template with caching support for improved performance."""
710 # Try to use environment's cache manager first, fall back to global cache
711 cache_manager = getattr(environment, "cache_manager", None)
712 if cache_manager:
713 # Use environment's cache manager
714 import hashlib
716 env_id = f"{id(environment)}:{getattr(environment, 'is_async', False)}"
717 content = f"{source}:{env_id}"
718 cache_key = hashlib.sha256(content.encode()).hexdigest()[:16]
720 # Check cache first
721 cached_code_first: str | None = cache_manager.get("compilation", cache_key)
722 if cached_code_first is not None:
723 return cached_code_first
725 # Compile and cache
726 generator = cls(environment, name, filename)
728 ast = environment.parse(source, name, filename)
729 compiled_code_first: str = generator.generate(ast)
731 # Store in cache
732 cache_manager.set("compilation", cache_key, compiled_code_first)
733 return compiled_code_first
734 else:
735 # Fall back to global cache for backward compatibility
736 env_id = f"{id(environment)}:{getattr(environment, 'is_async', False)}"
737 cache_key = _compilation_cache.get_cache_key(source, env_id)
739 # Check cache first
740 cached_code_second: str | None = _compilation_cache.get(cache_key)
741 if cached_code_second is not None:
742 return cached_code_second
744 # Compile and cache
745 generator = cls(environment, name, filename)
747 ast = environment.parse(source, name, filename)
748 compiled_code_second: str = generator.generate(ast)
750 # Store in cache
751 _compilation_cache.set(cache_key, compiled_code_second)
752 return compiled_code_second
754 def visit_For(self, node: nodes.For, frame: Frame) -> None:
755 frame = t.cast(AsyncFrame, frame)
756 if node.recursive:
757 raise NotImplementedError("Recursive loops not supported")
759 # Create frames and setup
760 loop_frame, test_frame, else_frame = self._setup_for_frames(frame)
761 extended_loop, loop_ref = self._setup_for_loop_context(node, loop_frame)
763 # Analyze nodes for variable declarations
764 self._analyze_for_nodes(node, loop_frame, else_frame)
766 # Handle loop filter
767 loop_filter_func = self._setup_for_filter(node, test_frame, loop_frame)
769 # Setup loop variables and checks
770 self._setup_for_variables(node, extended_loop, loop_ref)
772 # Generate main loop
773 iteration_indicator = self._generate_for_loop(
774 node, frame, loop_frame, extended_loop, loop_ref, loop_filter_func
775 )
777 # Handle else clause
778 self._handle_for_else(node, else_frame, iteration_indicator)
780 # Cleanup
781 self._cleanup_for_assignments(loop_frame)
783 def _setup_for_frames(
784 self, frame: AsyncFrame
785 ) -> tuple[AsyncFrame, AsyncFrame, AsyncFrame]:
786 """Setup frames for different scopes in for loop."""
787 loop_frame = frame.inner()
788 loop_frame.loop_frame = True
789 test_frame = frame.inner()
790 else_frame = frame.inner()
791 return loop_frame, test_frame, else_frame
793 def _setup_for_loop_context(
794 self, node: nodes.For, loop_frame: AsyncFrame
795 ) -> tuple[bool, str | None]:
796 """Setup extended loop context and loop reference."""
797 extended_loop = (
798 node.recursive
799 or "loop"
800 in find_undeclared(node.iter_child_nodes(only=("body",)), ("loop",))
801 or any(block.scoped for block in node.find_all(nodes.Block))
802 )
804 loop_ref = None
805 if extended_loop:
806 loop_ref = loop_frame.symbols.declare_parameter("loop")
808 return extended_loop, loop_ref
810 def _analyze_for_nodes(
811 self, node: nodes.For, loop_frame: AsyncFrame, else_frame: AsyncFrame
812 ) -> None:
813 """Analyze nodes for variable declarations."""
814 loop_frame.symbols.analyze_node(node, for_branch="body")
815 if node.else_:
816 else_frame.symbols.analyze_node(node, for_branch="else")
818 def _setup_for_filter(
819 self, node: nodes.For, test_frame: AsyncFrame, loop_frame: AsyncFrame
820 ) -> str | None:
821 """Setup loop filter if present."""
822 if not node.test:
823 return None
825 loop_filter_func = self.temporary_identifier()
826 test_frame.symbols.analyze_node(node, for_branch="test")
827 self.writeline(f"{self.func(loop_filter_func)}(filter):", node.test)
828 self.indent()
829 self.enter_frame(test_frame)
830 self.writeline(self.choose_async("async for ", "for "))
831 self.visit(node.target, loop_frame)
832 self.write(" in ")
833 self.write(self.choose_async("auto_aiter(filter)", "filter"))
834 self.write(":")
835 self.indent()
836 self.writeline("if ", node.test)
837 self.visit(node.test, test_frame)
838 self.write(":")
839 self.indent()
840 self.writeline("yield ")
841 self.visit(node.target, loop_frame)
842 self.outdent(3)
843 self.leave_frame(test_frame, with_python_scope=True)
844 return loop_filter_func
846 def _setup_for_variables(
847 self, node: nodes.For, extended_loop: bool, loop_ref: str | None
848 ) -> None:
849 """Setup loop variables and check for conflicts."""
850 if extended_loop and loop_ref:
851 self.writeline(f"{loop_ref} = missing")
853 for name in node.find_all(nodes.Name):
854 if name.ctx == "store" and name.name == "loop":
855 self.fail(
856 "Can't assign to special loop variable in for-loop target",
857 name.lineno,
858 )
860 def _generate_for_loop(
861 self,
862 node: nodes.For,
863 frame: AsyncFrame,
864 loop_frame: AsyncFrame,
865 extended_loop: bool,
866 loop_ref: str | None,
867 loop_filter_func: str | None,
868 ) -> str | None:
869 """Generate the main for loop code."""
870 # Handle else clause iteration indicator
871 iteration_indicator = None
872 if node.else_:
873 iteration_indicator = self.temporary_identifier()
874 self.writeline(f"{iteration_indicator} = 1")
876 # Generate the main loop using utility class
877 self._loop_generator.generate_async_for_header(node, node.target, loop_frame)
878 self._loop_generator.generate_loop_iterator(
879 node.iter, frame, extended_loop, loop_ref, loop_filter_func
880 )
882 self.indent()
883 self.enter_frame(loop_frame)
885 self.writeline("_loop_vars = {}")
886 self.blockvisit(node.body, loop_frame)
887 if node.else_:
888 self.writeline(f"{iteration_indicator} = 0")
889 self.outdent()
890 self.leave_frame(loop_frame, with_python_scope=not node.else_)
892 return iteration_indicator
894 def _handle_for_else(
895 self, node: nodes.For, else_frame: AsyncFrame, iteration_indicator: str | None
896 ) -> None:
897 """Handle the else clause of for loop."""
898 if not node.else_ or not iteration_indicator:
899 return
901 self.writeline(f"if {iteration_indicator}:")
902 self.indent()
903 self.enter_frame(else_frame)
904 self.blockvisit(node.else_, else_frame)
905 self.leave_frame(else_frame)
906 self.outdent()
908 def _cleanup_for_assignments(self, loop_frame: AsyncFrame) -> None:
909 """Clear assignments made in the loop from the top level."""
910 if hasattr(self, "_assign_stack") and self._assign_stack:
911 self._assign_stack[-1].difference_update(loop_frame.symbols.stores)
913 def visit_Macro(self, node: nodes.Macro, frame: Frame) -> None:
914 """Visit a macro node and generate async-aware code."""
915 frame = t.cast(AsyncFrame, frame)
916 # For now, let's just use the base class implementation without modification
917 # This ensures macros work in sync mode, and we can enhance async support later
918 super().visit_Macro(node, frame)
920 def visit_Filter(self, node: nodes.Filter, frame: Frame) -> None:
921 """Visit a filter node and generate async-aware code."""
922 frame = t.cast(AsyncFrame, frame)
924 filter_ref = self._get_filter_reference(node)
925 func = self.environment.filters.get(node.name)
927 if self.environment.is_async:
928 self.write("(await auto_await(")
930 self.write(f"{filter_ref}(")
931 self._write_filter_special_params(func)
932 self._write_filter_input(node, frame)
933 self._write_filter_arguments(node, frame)
934 self.write(")")
936 if self.environment.is_async:
937 self.write("))")
939 def _get_filter_reference(self, node: nodes.Filter) -> str:
940 """Get the filter reference from dependencies or fallback to environment."""
941 if node.name in self.filters:
942 return self.filters[node.name]
943 return f"environment.filters[{node.name!r}]"
945 def _write_filter_special_params(self, func: t.Any) -> None:
946 """Write special parameters that some filters need."""
947 # Try to get _PassArg using getattr to avoid import issues
948 _PassArg_module = None
949 try:
950 import jinja2.compiler
952 _PassArg_module = getattr(jinja2.compiler, "_PassArg", None)
953 except (ImportError, AttributeError):
954 _PassArg_module = None
956 pass_arg = None
957 if func and _PassArg_module is not None:
958 pass_arg_type = _PassArg_module.from_obj(func)
959 if pass_arg_type is not None:
960 # Use getattr to safely access attributes
961 context_attr = getattr(_PassArg_module, "context", None)
962 eval_context_attr = getattr(_PassArg_module, "eval_context", None)
963 environment_attr = getattr(_PassArg_module, "environment", None)
965 if context_attr is not None and pass_arg_type == context_attr:
966 pass_arg = "context"
967 elif (
968 eval_context_attr is not None and pass_arg_type == eval_context_attr
969 ):
970 pass_arg = "context.eval_ctx"
971 elif environment_attr is not None and pass_arg_type == environment_attr:
972 pass_arg = "environment"
974 if pass_arg is not None:
975 self.write(f"{pass_arg}, ")
977 def _write_filter_input(self, node: nodes.Filter, frame: AsyncFrame) -> None:
978 """Write the filter input value."""
979 if node.node is not None:
980 self.visit(node.node, frame)
981 elif frame.buffer is not None:
982 self._write_buffer_content(frame)
984 def _write_buffer_content(self, frame: AsyncFrame) -> None:
985 """Write buffer content for filter blocks."""
986 if frame.eval_ctx.volatile:
987 self.write(
988 f"(Markup(concat({frame.buffer}))"
989 f" if context.eval_ctx.autoescape else concat({frame.buffer}))"
990 )
991 elif frame.eval_ctx.autoescape:
992 self.write(f"Markup(concat({frame.buffer}))")
993 else:
994 self.write(f"concat({frame.buffer})")
996 def _write_filter_arguments(self, node: nodes.Filter, frame: AsyncFrame) -> None:
997 """Write filter arguments and keyword arguments."""
998 for arg in node.args:
999 self.write(", ")
1000 self.visit(arg, frame)
1002 for kwarg in node.kwargs:
1003 self.write(", ")
1004 self.visit(kwarg, frame)
1006 if node.dyn_args:
1007 self.write(", *")
1008 self.visit(node.dyn_args, frame)
1010 if node.dyn_kwargs:
1011 self.write(", **")
1012 self.visit(node.dyn_kwargs, frame)
1014 def visit_Assign(self, node: nodes.Assign, frame: Frame) -> None:
1015 """Visit an assignment node ({% set %} statements)."""
1016 frame = t.cast(AsyncFrame, frame)
1017 self.push_assign_tracking()
1019 # Check for namespace assignments like `ns.var = value`
1020 seen_refs: set[str] = set()
1021 for nsref in node.find_all(nodes.NSRef):
1022 if nsref.name in seen_refs:
1023 continue
1024 seen_refs.add(nsref.name)
1025 ref = frame.symbols.ref(nsref.name)
1026 self.writeline(f"if not isinstance({ref}, Namespace):")
1027 self.indent()
1028 self.writeline(
1029 "raise TemplateRuntimeError"
1030 '("cannot assign attribute on non-namespace object")'
1031 )
1032 self.outdent()
1034 # Generate the assignment code
1035 self.newline(node)
1036 self.visit(node.target, frame)
1037 self.write(" = ")
1038 self.visit(node.node, frame)
1039 self.pop_assign_tracking(frame)
1041 def push_assign_tracking(self) -> None:
1042 """Push a new layer for assignment tracking."""
1043 self._assign_stack.append(set())
1045 def pop_assign_tracking(self, frame: Frame) -> None:
1046 """Pop the topmost level for assignment tracking and update context variables."""
1047 frame = t.cast(AsyncFrame, frame)
1048 vars_set = self._assign_stack.pop()
1050 if (
1051 not frame.block_frame
1052 and not frame.loop_frame
1053 and not frame.toplevel
1054 or not vars_set
1055 ):
1056 return
1058 public_names = [x for x in vars_set if x[:1] != "_"]
1060 # Handle single variable case
1061 if len(vars_set) == 1:
1062 self._handle_single_variable(frame, vars_set)
1063 else:
1064 # Handle multiple variables case
1065 self._handle_multiple_variables(frame, vars_set)
1067 # Handle exported variables
1068 self._handle_exported_variables(frame, public_names)
1070 def _handle_single_variable(self, frame: AsyncFrame, vars_set: set[str]) -> None:
1071 """Handle the case with a single variable.
1073 Args:
1074 frame: Current frame
1075 vars_set: Set of variables
1076 """
1077 name = next(iter(vars_set))
1078 ref = frame.symbols.ref(name)
1079 if frame.loop_frame:
1080 self.writeline(f"_loop_vars[{name!r}] = {ref}")
1081 elif frame.block_frame:
1082 self.writeline(f"_block_vars[{name!r}] = {ref}")
1083 else:
1084 self.writeline(f"context.vars[{name!r}] = {ref}")
1086 def _handle_multiple_variables(self, frame: AsyncFrame, vars_set: set[str]) -> None:
1087 """Handle the case with multiple variables.
1089 Args:
1090 frame: Current frame
1091 vars_set: Set of variables
1092 """
1093 if frame.loop_frame:
1094 self.writeline("_loop_vars.update({")
1095 elif frame.block_frame:
1096 self.writeline("_block_vars.update({")
1097 else:
1098 self.writeline("context.vars.update({")
1100 for idx, name in enumerate(sorted(vars_set)):
1101 if idx:
1102 self.write(", ")
1103 ref = frame.symbols.ref(name)
1104 self.write(f"{name!r}: {ref}")
1105 self.write("})")
1107 def _handle_exported_variables(
1108 self, frame: AsyncFrame, public_names: list[str]
1109 ) -> None:
1110 """Handle exported variables.
1112 Args:
1113 frame: Current frame
1114 public_names: List of public variable names
1115 """
1116 if not frame.block_frame and not frame.loop_frame and public_names:
1117 if len(public_names) == 1:
1118 self.writeline(f"context.exported_vars.add({public_names[0]!r})")
1119 else:
1120 names_str = ", ".join(map(repr, sorted(public_names)))
1121 self.writeline(f"context.exported_vars.update(({names_str}))")