Coverage for jinja2_async_environment/compiler_modules/codegen.py: 81%

1"""Async code generator for template compilation."""

2

3 import typing as t

4

5 from jinja2 import nodes

6 from jinja2.compiler import (

7 CodeGenerator,

8 Frame,

9 find_undeclared,

10)

11

12 from .cache import CompilationCache

13 from .dependencies import DependencyResolver

14 from .frame import AsyncFrame

15 from .loops import LoopCodeGenerator

16 from .patterns import CompiledPatterns

17

18 if t.TYPE_CHECKING:

19 pass

20

21# Global compilation cache instance

22 _compilation_cache = CompilationCache()

23

24

25 class AsyncCodeGenerator(CodeGenerator):

26 """Async-aware code generator extending Jinja2's CodeGenerator."""

27

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] = []

45

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 )

61

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 = []

79

80 # Initialize assignment tracking stack

81 self._assign_stack: list[set[str]] = []

82

83 # Initialize utility classes for better code organization

84 self._dependency_resolver = DependencyResolver(self)

85 self._loop_generator = LoopCodeGenerator(self)

86

87 from jinja2.nodes import EvalContext

88

89 if self.eval_ctx is None:

90 self.eval_ctx = EvalContext(self.environment, self.name)

91

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

94

95 def simple_write(self, value: str, frame: Frame) -> None: # type: ignore[override]

96 self.writeline(f"yield {value}")

97

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"

101

102 def func(self, name: str) -> str:

103 """Generate a function declaration for the given name.

104

105 Properly handles async functions by checking environment.is_async.

106 """

107 return f"{self.choose_async()}def {name}"

108

109 def enter_frame(self, frame: Frame) -> None:

110 """Enter a new frame context and load variables from context.

111

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 )

121

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")

136

137 def leave_frame(self, frame: Frame, with_python_scope: bool = False) -> None:

138 """Leave a frame context."""

139 pass

140

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})")

149

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)

156

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)

163

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

169

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

184

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}]")

201

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

206

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

211

212 if VAR_LOAD_PARAMETER is None:

213 return False

214

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 )

222

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

226

227 visitor = DependencyFinderVisitor()

228 for node in nodes:

229 visitor.visit(node)

230

231 # Set up filter dependencies using utility class

232 for name in sorted(visitor.filters):

233 self._dependency_resolver.setup_filter_dependency(name)

234

235 # Set up test dependencies using utility class

236 for name in sorted(visitor.tests):

237 self._dependency_resolver.setup_test_dependency(name)

238

239 def generate(self, node: nodes.Template) -> str:

240 """Generate template code following base Jinja2 architecture.

241

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

248

249 This separation of block calling vs. definition is CRITICAL for

250 template inheritance to work correctly.

251 """

252 from jinja2.nodes import EvalContext

253

254 if self.eval_ctx is None:

255 self.eval_ctx = EvalContext(self.environment, self.name)

256

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

263

264 # Check if template has extends (for proper frame configuration)

265 have_extends = node.find(nodes.Extends) is not None

266

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)

271

272 self.writeline(f"name = {self.name! r}")

273

274 # Helper functions

275 self.writeline("def undefined(name=None, **_):")

276 self.indent()

277 self.writeline("return Undefined(name=name)")

278 self.outdent()

279

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()

288

289 self.writeline("filters = DEFAULT_FILTERS.copy()")

290 self.writeline("filters['escape'] = escape")

291

292 # PHASE 3: Generate root render function

293 self.writeline("async def root(context):")

294 self.indent()

295

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

303

304 # Create frame for root processing

305 frame = self.root_frame_class(eval_ctx=self.eval_ctx)

306

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)")

311

312 # Analyze node to discover variables

313 frame.symbols.analyze_node(node)

314

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

318

319 # Initialize parent_template for extends

320 if have_extends:

321 self.writeline("parent_template = None")

322

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

330

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))

348

349 # PHASE 5: Generate all block function DEFINITIONS

350 # This happens AFTER root function exits

351 self._generate_block_functions()

352

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")

357

358 # Apply pattern-based optimizations to generated code

359 generated_code = self.stream.getvalue()

360 return CompiledPatterns.optimize_generated_code(generated_code)

361

362 def _generate_block_functions(self) -> None:

363 """Generate all block function definitions.

364

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.

368

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()

383

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

390

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

394

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})")

403

404 # Analyze block body

405 block_frame.symbols.analyze_node(block_node)

406 block_frame.block = name

407

408 # Block-level variables dict

409 self.writeline("_block_vars = {}")

410

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)

416

417 self.outdent() # Exit block function

418

419 def visit_Block(self, node: nodes.Block, frame: Frame) -> None:

420 """Call a block and register it for the template.

421

422 This method ONLY handles CALLING blocks from context.blocks.

423 Block DEFINITIONS are generated separately in _generate_block_functions().

424

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]

429

430 Architecture matches base Jinja2's CodeGenerator.visit_Block.

431 """

432 block_name = node.name

433

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

445

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()

451

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()

461

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()")

473

474 # Close inheritance conditional if needed

475 self.outdent(level)

476

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

480

481 # If output check is not required, raise CompilerExit immediately

482 if not frame.require_output_check:

483 from jinja2.compiler import CompilerExit

484

485 raise CompilerExit()

486

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)

490

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

499

500 raise CompilerExit()

501 else:

502 self.outdent()

503

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()

512

513 # Update inheritance tracking

514 if frame.rootlevel:

515 self.has_known_extends = True

516 self.extends_so_far += 1

517

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

521

522 # Handle ignore_missing flag

523 if node.ignore_missing:

524 self.writeline("try:")

525 self.indent()

526

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})")

531

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()

541

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 )

554

555 # Generate event output

556 self.indent()

557 self.simple_write("event", frame)

558 self.outdent()

559

560 # Close else block for ignore_missing

561 if node.ignore_missing:

562 self.outdent()

563

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

567

568 # Handle recursive loops (not supported)

569 if node.recursive:

570 raise NotImplementedError("Recursive loops not supported")

571

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)

579

580 # Initialize target variable

581 self.writeline(f"{item} = None")

582

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)

589

590 # Initialize loop counter

591 loop_var = self.temporary_identifier()

592 self.writeline(f"{loop_var} = -1", node)

593

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()

599

600 # Increment loop counter

601 self.writeline(f"{loop_var} += 1")

602

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()

607

608 # Process loop body

609 if hasattr(node, "body"):

610 self.blockvisit(node.body, frame)

611

612 # Close filter condition

613 if hasattr(node, "test") and node.test and loop_filter:

614 self.outdent()

615

616 # Close main loop

617 self.outdent()

618

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()

625

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)

630

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

634

635 # Early return if no filter or body

636 if not hasattr(node, "filter"):

637 return

638 if not hasattr(node, "body"):

639 return

640

641 # Get filter node

642 filter_node = node.filter

643

644 # Create buffer for collecting content

645 buffer = self.temporary_identifier()

646 self.writeline(f"{buffer} = []")

647

648 # Create async frame for processing body

649 asyncframe = frame.copy()

650 asyncframe.buffer = buffer

651 asyncframe.toplevel = False

652

653 # Process the body

654 self.blockvisit(node.body, asyncframe)

655

656 # Generate await call for filter

657 self.writeline("await ", filter_node)

658 self.visit(filter_node, frame)

659 self.write(f"(''.join({buffer}))")

660

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

664

665 # Early return if no name or body

666 if not hasattr(node, "name"):

667 return

668 if not hasattr(node, "body"):

669 return

670

671 # Get block name

672 block_name = node.name

673

674 # Initialize block storage

675 self.writeline(f"blocks[{block_name! r}] = []")

676

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()

681

682 # Empty block content placeholder

683 self.writeline("yield ''")

684

685 # Process block body if it exists

686 if node.body:

687 self.blockvisit(node.body, frame)

688

689 # End function and register block

690 self.outdent()

691 self.writeline(f"blocks[{block_name! r}].append({block_func_name})")

692

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)

699

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})")

704

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

715

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]

719

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

724

725 # Compile and cache

726 generator = cls(environment, name, filename)

727

728 ast = environment.parse(source, name, filename)

729 compiled_code_first: str = generator.generate(ast)

730

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)

738

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

743

744 # Compile and cache

745 generator = cls(environment, name, filename)

746

747 ast = environment.parse(source, name, filename)

748 compiled_code_second: str = generator.generate(ast)

749

750 # Store in cache

751 _compilation_cache.set(cache_key, compiled_code_second)

752 return compiled_code_second

753

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")

758

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)

762

763 # Analyze nodes for variable declarations

764 self._analyze_for_nodes(node, loop_frame, else_frame)

765

766 # Handle loop filter

767 loop_filter_func = self._setup_for_filter(node, test_frame, loop_frame)

768

769 # Setup loop variables and checks

770 self._setup_for_variables(node, extended_loop, loop_ref)

771

772 # Generate main loop

773 iteration_indicator = self._generate_for_loop(

774 node, frame, loop_frame, extended_loop, loop_ref, loop_filter_func

775 )

776

777 # Handle else clause

778 self._handle_for_else(node, else_frame, iteration_indicator)

779

780 # Cleanup

781 self._cleanup_for_assignments(loop_frame)

782

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

792

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 )

803

804 loop_ref = None

805 if extended_loop:

806 loop_ref = loop_frame.symbols.declare_parameter("loop")

807

808 return extended_loop, loop_ref

809

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")

817

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

824

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

845

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")

852

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 )

859

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")

875

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 )

881

882 self.indent()

883 self.enter_frame(loop_frame)

884

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_)

891

892 return iteration_indicator

893

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

900

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()

907

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)

912

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)

919

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)

923

924 filter_ref = self._get_filter_reference(node)

925 func = self.environment.filters.get(node.name)

926

927 if self.environment.is_async:

928 self.write("(await auto_await(")

929

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(")")

935

936 if self.environment.is_async:

937 self.write("))")

938

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}]"

944

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

951

952 _PassArg_module = getattr(jinja2.compiler, "_PassArg", None)

953 except (ImportError, AttributeError):

954 _PassArg_module = None

955

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)

964

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"

973

974 if pass_arg is not None:

975 self.write(f"{pass_arg}, ")

976

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)

983

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})")

995

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)

1001

1002 for kwarg in node.kwargs:

1003 self.write(", ")

1004 self.visit(kwarg, frame)

1005

1006 if node.dyn_args:

1007 self.write(", *")

1008 self.visit(node.dyn_args, frame)

1009

1010 if node.dyn_kwargs:

1011 self.write(", **")

1012 self.visit(node.dyn_kwargs, frame)

1013

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()

1018

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()

1033

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)

1040

1041 def push_assign_tracking(self) -> None:

1042 """Push a new layer for assignment tracking."""

1043 self._assign_stack.append(set())

1044

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()

1049

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

1057

1058 public_names = [x for x in vars_set if x[:1] != "_"]

1059

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)

1066

1067 # Handle exported variables

1068 self._handle_exported_variables(frame, public_names)

1069

1070 def _handle_single_variable(self, frame: AsyncFrame, vars_set: set[str]) -> None:

1071 """Handle the case with a single variable.

1072

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}")

1085

1086 def _handle_multiple_variables(self, frame: AsyncFrame, vars_set: set[str]) -> None:

1087 """Handle the case with multiple variables.

1088

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({")

1099

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("})")

1106

1107 def _handle_exported_variables(

1108 self, frame: AsyncFrame, public_names: list[str]

1109 ) -> None:

1110 """Handle exported variables.

1111

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}))")