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

512 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-09-03 14:09 -0700

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

2 

3import typing as t 

4 

5from jinja2 import nodes 

6from jinja2.compiler import ( 

7 CodeGenerator, 

8 Frame, 

9 find_undeclared, 

10) 

11 

12from .cache import CompilationCache 

13from .dependencies import DependencyResolver 

14from .frame import AsyncFrame 

15from .loops import LoopCodeGenerator 

16from .patterns import CompiledPatterns 

17 

18if t.TYPE_CHECKING: 

19 pass 

20 

21# Global compilation cache instance 

22_compilation_cache = CompilationCache() 

23 

24 

25class 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 return f"def {name}" 

105 

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

107 """Enter a new frame context.""" 

108 pass 

109 

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

111 """Leave a frame context.""" 

112 pass 

113 

114 def return_buffer_contents( 

115 self, 

116 frame: Frame, 

117 force_unescaped: bool = False, # noqa: ARG002 

118 ) -> None: 

119 _ = force_unescaped 

120 if frame.buffer is not None: 

121 self.writeline(f"return ''.join({frame.buffer})") 

122 

123 def visit_Name(self, node: nodes.Name, frame: Frame) -> None: 

124 frame = t.cast(AsyncFrame, frame) 

125 self._handle_assignment_tracking(node, frame) 

126 if self._handle_special_names(node): 

127 return 

128 self._handle_symbol_name(node, frame) 

129 

130 def _handle_assignment_tracking(self, node: nodes.Name, frame: AsyncFrame) -> None: 

131 if node.ctx == "store": 

132 frame.symbols.store(node.name) 

133 if frame.toplevel or frame.loop_frame or frame.block_frame: 

134 if hasattr(self, "_assign_stack") and self._assign_stack: 

135 self._assign_stack[-1].add(node.name) 

136 

137 def _handle_special_names(self, node: nodes.Name) -> bool: 

138 if node.name in ("blocks", "debug_info"): 

139 self.write(node.name) 

140 return True 

141 return False 

142 

143 def _handle_symbol_name(self, node: nodes.Name, frame: AsyncFrame) -> None: 

144 # Fast path for common variables 

145 if node.name in self._COMMON_VARS and node.ctx == "load": 

146 try: 

147 ref = frame.symbols.ref(node.name) 

148 self.write(ref) 

149 return 

150 except AssertionError: 

151 self.write(f"context.get({node.name!r})") 

152 return 

153 

154 # Standard path for other variables 

155 try: 

156 ref = frame.symbols.ref(node.name) 

157 if node.ctx == "load" and self._should_use_undefined_check(ref, frame): 

158 self.write( 

159 f"(undefined(name={node.name!r}) if {ref} is missing else {ref})" 

160 ) 

161 else: 

162 self.write(ref) 

163 except AssertionError: 

164 if node.ctx == "load": 

165 self.write(f"context.get({node.name!r})") 

166 else: 

167 self.write(f"context.vars[{node.name!r}]") 

168 

169 def _should_use_undefined_check(self, ref: str, frame: AsyncFrame) -> bool: 

170 from jinja2.compiler import VAR_LOAD_PARAMETER 

171 

172 load = frame.symbols.find_load(ref) 

173 return not ( 

174 load is not None 

175 and load[0] == VAR_LOAD_PARAMETER 

176 and hasattr(self, "parameter_is_undeclared") 

177 and not self.parameter_is_undeclared(ref) 

178 ) 

179 

180 def pull_dependencies(self, nodes: t.Iterable[nodes.Node]) -> None: 

181 """Find all filter and test names used in the template and assign them to variables.""" 

182 from jinja2.compiler import DependencyFinderVisitor 

183 

184 visitor = DependencyFinderVisitor() 

185 for node in nodes: 

186 visitor.visit(node) 

187 

188 # Set up filter dependencies using utility class 

189 for name in sorted(visitor.filters): 

190 self._dependency_resolver.setup_filter_dependency(name) 

191 

192 # Set up test dependencies using utility class 

193 for name in sorted(visitor.tests): 

194 self._dependency_resolver.setup_test_dependency(name) 

195 

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

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

198 self.writeline("blocks = {}") 

199 self.writeline("debug_info = None") 

200 

201 # Use optimized cached imports for better performance 

202 for import_line in CompiledPatterns.get_optimized_imports().split("\n"): 

203 self.writeline(import_line) 

204 

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

206 self.indent() 

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

208 self.outdent() 

209 

210 self.writeline("async def auto_await(value):") 

211 self.indent() 

212 self.writeline("if hasattr(value, '__await__'):") 

213 self.indent() 

214 self.writeline("return await value") 

215 self.outdent() 

216 self.writeline("return value") 

217 self.outdent() 

218 

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

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

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

222 self.indent() 

223 self.writeline("parent_template = None") 

224 self.writeline("environment = context.environment") 

225 self.writeline("eval_ctx = context.eval_ctx") 

226 self.writeline("undefined = environment.undefined") 

227 

228 from jinja2.nodes import EvalContext 

229 

230 if self.eval_ctx is None: 

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

232 

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

234 frame.toplevel = frame.rootlevel = True 

235 frame.require_output_check = False 

236 frame.buffer = None 

237 

238 for macro in node.find_all(nodes.Macro): 

239 frame.symbols.store(macro.name) 

240 

241 # Pull dependencies for filters and tests 

242 self.pull_dependencies(node.body) 

243 self.blockvisit(node.body, frame) 

244 self.outdent() 

245 

246 # Apply pattern-based optimizations to generated code 

247 generated_code = self.stream.getvalue() 

248 return CompiledPatterns.optimize_generated_code(generated_code) 

249 

250 def visit_Block(self, node: nodes.Block, frame: AsyncFrame) -> None: 

251 """Visit a block node with proper async handling.""" 

252 # Cast frame to AsyncFrame for type safety 

253 frame = t.cast(AsyncFrame, frame) 

254 

255 # Get block name 

256 block_name = node.name 

257 

258 # Initialize block storage 

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

260 

261 # Define block function with async prefix 

262 block_func_name = f"block_{block_name}" 

263 async_prefix = self.choose_async() 

264 self.writeline(f"{async_prefix}def {block_func_name}(context):") 

265 self.indent() 

266 

267 # Empty block content placeholder 

268 self.writeline("yield ''") 

269 

270 # Process block body if it exists 

271 if node.body: 

272 self.blockvisit(node.body, frame) 

273 

274 # End function and register block 

275 self.outdent() 

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

277 

278 # Handle inheritance cases 

279 level = 0 

280 if frame.toplevel: 

281 if self.has_known_extends: 

282 return 

283 if self.extends_so_far > 0: 

284 self.writeline("if parent_template is None:") 

285 self.indent() 

286 level += 1 

287 

288 # Handle scoped blocks 

289 if node.scoped: 

290 context = self.derive_context(frame) 

291 else: 

292 context = self.get_context_ref() 

293 

294 # Generate block call with proper async handling 

295 self.writeline(f"{async_prefix}for event in {block_func_name}({context}):") 

296 self.indent() 

297 self.writeline("yield event") 

298 self.outdent() 

299 

300 # Close conditional blocks 

301 for _ in range(level): 

302 self.outdent() 

303 

304 def visit_Extends(self, node: nodes.Extends, frame: AsyncFrame) -> None: 

305 """Visit an extends node with proper async handling.""" 

306 # Cast frame to AsyncFrame for type safety 

307 frame = t.cast(AsyncFrame, frame) 

308 

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

310 if not frame.require_output_check: 

311 from jinja2.compiler import CompilerExit 

312 

313 raise CompilerExit() 

314 

315 # Check if we're in a top-level scope 

316 if not frame.toplevel: 

317 self.fail("cannot use extend from a non top-level scope", node.lineno) 

318 

319 # Handle multiple extends 

320 if self.extends_so_far > 0: 

321 if not self.has_known_extends: 

322 self.writeline("if parent_template is not None:") 

323 self.indent() 

324 self.writeline('raise TemplateRuntimeError("extended multiple times")') 

325 if self.has_known_extends: 

326 from jinja2.compiler import CompilerExit 

327 

328 raise CompilerExit() 

329 else: 

330 self.outdent() 

331 

332 # Generate async template loading code 

333 self.writeline("parent_template = await environment.get_template_async(", node) 

334 self.visit(node.template, frame) 

335 self.write(f", {self.name!r})") 

336 self.writeline("for name, parent_block in parent_template.blocks.items():") 

337 self.indent() 

338 self.writeline("context.blocks.setdefault(name, []).append(parent_block)") 

339 self.outdent() 

340 

341 # Update inheritance tracking 

342 if frame.rootlevel: 

343 self.has_known_extends = True 

344 self.extends_so_far += 1 

345 

346 def visit_Include(self, node: nodes.Include, frame: AsyncFrame) -> None: 

347 """Visit an include node with proper async handling.""" 

348 # Cast frame to AsyncFrame for type safety 

349 frame = t.cast(AsyncFrame, frame) 

350 

351 # Handle ignore_missing flag 

352 if node.ignore_missing: 

353 self.writeline("try:") 

354 self.indent() 

355 

356 # Generate async template loading code 

357 self.writeline("template = await environment.get_template_async(", node) 

358 self.visit(node.template, frame) 

359 self.write(f", {self.name!r})") 

360 

361 # Close try block for ignore_missing 

362 if node.ignore_missing: 

363 self.outdent() 

364 self.writeline("except TemplateNotFound:") 

365 self.indent() 

366 self.writeline("pass") 

367 self.outdent() 

368 self.writeline("else:") 

369 self.indent() 

370 

371 # Generate rendering code based on context flag 

372 if node.with_context: 

373 # With context - include local variables 

374 local_context = self.dump_local_context(frame) 

375 self.writeline( 

376 f"async for event in template.root_render_func(template.new_context(context.get_all(), True, {local_context})):" 

377 ) 

378 else: 

379 # Without context - use default module 

380 self.writeline( 

381 "async for event in (await template._get_default_module_async())._body_stream:" 

382 ) 

383 

384 # Generate event output 

385 self.indent() 

386 self.simple_write("event", frame) 

387 self.outdent() 

388 

389 # Close else block for ignore_missing 

390 if node.ignore_missing: 

391 self.outdent() 

392 

393 def visit_AsyncFor(self, node: nodes.Node, frame: AsyncFrame) -> None: 

394 """Visit an async for loop node with proper async handling.""" 

395 # Cast frame to AsyncFrame for type safety 

396 frame = t.cast(AsyncFrame, frame) 

397 

398 # Handle recursive loops (not supported) 

399 if hasattr(node, "recursive") and node.recursive: 

400 raise NotImplementedError("Recursive loops not supported") 

401 

402 # Get target variable name 

403 target = node.target 

404 item = target.name if hasattr(target, "name") else "item" 

405 frame.symbols.store(item) 

406 

407 # Initialize target variable 

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

409 

410 # Handle loop filter 

411 loop_filter = None 

412 if hasattr(node, "test") and node.test: 

413 loop_filter = self.temporary_identifier() 

414 self.writeline(f"{loop_filter} = ", node.test) 

415 self.visit(node.test, frame) 

416 

417 # Initialize loop counter 

418 loop_var = self.temporary_identifier() 

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

420 

421 # Generate async for loop 

422 self.writeline(f"async for {item} in ", node.iter) 

423 self.visit(node.iter, frame) 

424 self.write(":") 

425 self.indent() 

426 

427 # Increment loop counter 

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

429 

430 # Handle loop filter condition 

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

432 self.writeline(f"if {loop_filter}({item}):") 

433 self.indent() 

434 

435 # Process loop body 

436 if hasattr(node, "body"): 

437 self.blockvisit(node.body, frame) 

438 

439 # Close filter condition 

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

441 self.outdent() 

442 

443 # Close main loop 

444 self.outdent() 

445 

446 # Handle else clause 

447 if hasattr(node, "else_") and node.else_: 

448 self.writeline(f"if {loop_var} == -1:") 

449 self.indent() 

450 self.blockvisit(node.else_, frame) 

451 self.outdent() 

452 

453 def visit_AsyncCall(self, node: nodes.Node, frame: AsyncFrame) -> None: 

454 """Visit an async call node by adding await prefix.""" 

455 self.write("await ") 

456 self.visit_Call(node, frame) 

457 

458 def visit_AsyncFilterBlock(self, node: nodes.Node, frame: AsyncFrame) -> None: 

459 """Visit an async filter block node.""" 

460 # Cast frame to AsyncFrame for type safety 

461 frame = t.cast(AsyncFrame, frame) 

462 

463 # Early return if no filter or body 

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

465 return 

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

467 return 

468 

469 # Get filter node 

470 filter_node = node.filter 

471 

472 # Create buffer for collecting content 

473 buffer = self.temporary_identifier() 

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

475 

476 # Create async frame for processing body 

477 asyncframe = frame.copy() 

478 asyncframe.buffer = buffer 

479 asyncframe.toplevel = False 

480 

481 # Process the body 

482 self.blockvisit(node.body, asyncframe) 

483 

484 # Generate await call for filter 

485 self.writeline("await ", filter_node) 

486 self.visit(filter_node, frame) 

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

488 

489 def visit_AsyncBlock(self, node: nodes.Node, frame: AsyncFrame) -> None: 

490 """Visit an async block node.""" 

491 # Cast frame to AsyncFrame for type safety 

492 frame = t.cast(AsyncFrame, frame) 

493 

494 # Early return if no name or body 

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

496 return 

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

498 return 

499 

500 # Get block name 

501 block_name = node.name 

502 

503 # Initialize block storage 

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

505 

506 # Define async block function 

507 block_func_name = f"block_{block_name}" 

508 self.writeline(f"async def {block_func_name}(context):") 

509 self.indent() 

510 

511 # Empty block content placeholder 

512 self.writeline("yield ''") 

513 

514 # Process block body if it exists 

515 if node.body: 

516 self.blockvisit(node.body, frame) 

517 

518 # End function and register block 

519 self.outdent() 

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

521 

522 def _import_common( 

523 self, node: nodes.Import | nodes.FromImport, frame: AsyncFrame 

524 ) -> None: 

525 """Common import functionality with async template loading.""" 

526 # Cast frame to AsyncFrame for type safety 

527 frame = t.cast(AsyncFrame, frame) 

528 

529 # Generate async template loading code 

530 self.writeline("template = await environment.get_template_async(", node) 

531 self.visit(node.template, frame) 

532 self.write(f", {self.name!r})") 

533 

534 @classmethod 

535 def compile_with_cache( 

536 cls, environment: t.Any, source: str, name: str, filename: str 

537 ) -> str: 

538 """Compile template with caching support for improved performance.""" 

539 # Try to use environment's cache manager first, fall back to global cache 

540 cache_manager = getattr(environment, "cache_manager", None) 

541 if cache_manager: 

542 # Use environment's cache manager 

543 import hashlib 

544 

545 env_id = f"{id(environment)}:{getattr(environment, 'is_async', False)}" 

546 content = f"{source}:{env_id}" 

547 cache_key = hashlib.sha256(content.encode()).hexdigest()[:16] 

548 

549 # Check cache first 

550 cached_code = cache_manager.get("compilation", cache_key) 

551 if cached_code is not None: 

552 return cached_code 

553 

554 # Compile and cache 

555 generator = cls(environment, name, filename) 

556 

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

558 compiled_code = generator.generate(ast) 

559 

560 # Store in cache 

561 cache_manager.set("compilation", cache_key, compiled_code) 

562 return compiled_code 

563 else: 

564 # Fall back to global cache for backward compatibility 

565 env_id = f"{id(environment)}:{getattr(environment, 'is_async', False)}" 

566 cache_key = _compilation_cache.get_cache_key(source, env_id) 

567 

568 # Check cache first 

569 cached_code = _compilation_cache.get(cache_key) 

570 if cached_code is not None: 

571 return cached_code 

572 

573 # Compile and cache 

574 generator = cls(environment, name, filename) 

575 

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

577 compiled_code = generator.generate(ast) 

578 

579 # Store in cache 

580 _compilation_cache.set(cache_key, compiled_code) 

581 return compiled_code 

582 

583 def visit_For(self, node: nodes.For, frame: Frame) -> None: 

584 frame = t.cast(AsyncFrame, frame) 

585 if node.recursive: 

586 raise NotImplementedError("Recursive loops not supported") 

587 

588 # Create frames and setup 

589 loop_frame, test_frame, else_frame = self._setup_for_frames(frame) 

590 extended_loop, loop_ref = self._setup_for_loop_context(node, loop_frame) 

591 

592 # Analyze nodes for variable declarations 

593 self._analyze_for_nodes(node, loop_frame, else_frame) 

594 

595 # Handle loop filter 

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

597 

598 # Setup loop variables and checks 

599 self._setup_for_variables(node, extended_loop, loop_ref) 

600 

601 # Generate main loop 

602 iteration_indicator = self._generate_for_loop( 

603 node, frame, loop_frame, extended_loop, loop_ref, loop_filter_func 

604 ) 

605 

606 # Handle else clause 

607 self._handle_for_else(node, else_frame, iteration_indicator) 

608 

609 # Cleanup 

610 self._cleanup_for_assignments(loop_frame) 

611 

612 def _setup_for_frames( 

613 self, frame: AsyncFrame 

614 ) -> tuple[AsyncFrame, AsyncFrame, AsyncFrame]: 

615 """Setup frames for different scopes in for loop.""" 

616 loop_frame = frame.inner() 

617 loop_frame.loop_frame = True 

618 test_frame = frame.inner() 

619 else_frame = frame.inner() 

620 return loop_frame, test_frame, else_frame 

621 

622 def _setup_for_loop_context( 

623 self, node: nodes.For, loop_frame: AsyncFrame 

624 ) -> tuple[bool, str | None]: 

625 """Setup extended loop context and loop reference.""" 

626 extended_loop = ( 

627 node.recursive 

628 or "loop" 

629 in find_undeclared(node.iter_child_nodes(only=("body",)), ("loop",)) 

630 or any(block.scoped for block in node.find_all(nodes.Block)) 

631 ) 

632 

633 loop_ref = None 

634 if extended_loop: 

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

636 

637 return extended_loop, loop_ref 

638 

639 def _analyze_for_nodes( 

640 self, node: nodes.For, loop_frame: AsyncFrame, else_frame: AsyncFrame 

641 ) -> None: 

642 """Analyze nodes for variable declarations.""" 

643 loop_frame.symbols.analyze_node(node, for_branch="body") 

644 if node.else_: 

645 else_frame.symbols.analyze_node(node, for_branch="else") 

646 

647 def _setup_for_filter( 

648 self, node: nodes.For, test_frame: AsyncFrame, loop_frame: AsyncFrame 

649 ) -> str | None: 

650 """Setup loop filter if present.""" 

651 if not node.test: 

652 return None 

653 

654 loop_filter_func = self.temporary_identifier() 

655 test_frame.symbols.analyze_node(node, for_branch="test") 

656 self.writeline(f"{self.func(loop_filter_func)}(filter):", node.test) 

657 self.indent() 

658 self.enter_frame(test_frame) 

659 self.writeline(self.choose_async("async for ", "for ")) 

660 self.visit(node.target, loop_frame) 

661 self.write(" in ") 

662 self.write(self.choose_async("auto_aiter(filter)", "filter")) 

663 self.write(":") 

664 self.indent() 

665 self.writeline("if ", node.test) 

666 self.visit(node.test, test_frame) 

667 self.write(":") 

668 self.indent() 

669 self.writeline("yield ") 

670 self.visit(node.target, loop_frame) 

671 self.outdent(3) 

672 self.leave_frame(test_frame, with_python_scope=True) 

673 return loop_filter_func 

674 

675 def _setup_for_variables( 

676 self, node: nodes.For, extended_loop: bool, loop_ref: str | None 

677 ) -> None: 

678 """Setup loop variables and check for conflicts.""" 

679 if extended_loop and loop_ref: 

680 self.writeline(f"{loop_ref} = missing") 

681 

682 for name in node.find_all(nodes.Name): 

683 if name.ctx == "store" and name.name == "loop": 

684 self.fail( 

685 "Can't assign to special loop variable in for-loop target", 

686 name.lineno, 

687 ) 

688 

689 def _generate_for_loop( 

690 self, 

691 node: nodes.For, 

692 frame: AsyncFrame, 

693 loop_frame: AsyncFrame, 

694 extended_loop: bool, 

695 loop_ref: str | None, 

696 loop_filter_func: str | None, 

697 ) -> str | None: 

698 """Generate the main for loop code.""" 

699 # Handle else clause iteration indicator 

700 iteration_indicator = None 

701 if node.else_: 

702 iteration_indicator = self.temporary_identifier() 

703 self.writeline(f"{iteration_indicator} = 1") 

704 

705 # Generate the main loop using utility class 

706 self._loop_generator.generate_async_for_header(node, node.target, loop_frame) 

707 self._loop_generator.generate_loop_iterator( 

708 node.iter, frame, extended_loop, loop_ref, loop_filter_func 

709 ) 

710 

711 self.indent() 

712 self.enter_frame(loop_frame) 

713 

714 self.writeline("_loop_vars = {}") 

715 self.blockvisit(node.body, loop_frame) 

716 if node.else_: 

717 self.writeline(f"{iteration_indicator} = 0") 

718 self.outdent() 

719 self.leave_frame(loop_frame, with_python_scope=not node.else_) 

720 

721 return iteration_indicator 

722 

723 def _handle_for_else( 

724 self, node: nodes.For, else_frame: AsyncFrame, iteration_indicator: str | None 

725 ) -> None: 

726 """Handle the else clause of for loop.""" 

727 if not node.else_ or not iteration_indicator: 

728 return 

729 

730 self.writeline(f"if {iteration_indicator}:") 

731 self.indent() 

732 self.enter_frame(else_frame) 

733 self.blockvisit(node.else_, else_frame) 

734 self.leave_frame(else_frame) 

735 self.outdent() 

736 

737 def _cleanup_for_assignments(self, loop_frame: AsyncFrame) -> None: 

738 """Clear assignments made in the loop from the top level.""" 

739 if hasattr(self, "_assign_stack") and self._assign_stack: 

740 self._assign_stack[-1].difference_update(loop_frame.symbols.stores) 

741 

742 def visit_Macro(self, node: nodes.Macro, frame: Frame) -> None: 

743 """Visit a macro node and generate async-aware code.""" 

744 frame = t.cast(AsyncFrame, frame) 

745 # For now, let's just use the base class implementation without modification 

746 # This ensures macros work in sync mode, and we can enhance async support later 

747 super().visit_Macro(node, frame) 

748 

749 def visit_Filter(self, node: nodes.Filter, frame: Frame) -> None: 

750 """Visit a filter node and generate async-aware code.""" 

751 frame = t.cast(AsyncFrame, frame) 

752 

753 filter_ref = self._get_filter_reference(node) 

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

755 

756 if self.environment.is_async: 

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

758 

759 self.write(f"{filter_ref}(") 

760 self._write_filter_special_params(func) 

761 self._write_filter_input(node, frame) 

762 self._write_filter_arguments(node, frame) 

763 self.write(")") 

764 

765 if self.environment.is_async: 

766 self.write("))") 

767 

768 def _get_filter_reference(self, node: nodes.Filter) -> str: 

769 """Get the filter reference from dependencies or fallback to environment.""" 

770 if node.name in self.filters: 

771 return self.filters[node.name] 

772 return f"environment.filters[{node.name!r}]" 

773 

774 def _write_filter_special_params(self, func: t.Any) -> None: 

775 """Write special parameters that some filters need.""" 

776 from jinja2.compiler import _PassArg 

777 

778 pass_arg = None 

779 if func: 

780 pass_arg_type = _PassArg.from_obj(func) 

781 if pass_arg_type: 

782 pass_arg = { 

783 _PassArg.context: "context", 

784 _PassArg.eval_context: "context.eval_ctx", 

785 _PassArg.environment: "environment", 

786 }.get(pass_arg_type) 

787 

788 if pass_arg is not None: 

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

790 

791 def _write_filter_input(self, node: nodes.Filter, frame: AsyncFrame) -> None: 

792 """Write the filter input value.""" 

793 if node.node is not None: 

794 self.visit(node.node, frame) 

795 elif frame.buffer is not None: 

796 self._write_buffer_content(frame) 

797 

798 def _write_buffer_content(self, frame: AsyncFrame) -> None: 

799 """Write buffer content for filter blocks.""" 

800 if frame.eval_ctx.volatile: 

801 self.write( 

802 f"(Markup(concat({frame.buffer}))" 

803 f" if context.eval_ctx.autoescape else concat({frame.buffer}))" 

804 ) 

805 elif frame.eval_ctx.autoescape: 

806 self.write(f"Markup(concat({frame.buffer}))") 

807 else: 

808 self.write(f"concat({frame.buffer})") 

809 

810 def _write_filter_arguments(self, node: nodes.Filter, frame: AsyncFrame) -> None: 

811 """Write filter arguments and keyword arguments.""" 

812 for arg in node.args: 

813 self.write(", ") 

814 self.visit(arg, frame) 

815 

816 for kwarg in node.kwargs: 

817 self.write(", ") 

818 self.visit(kwarg, frame) 

819 

820 if node.dyn_args: 

821 self.write(", *") 

822 self.visit(node.dyn_args, frame) 

823 

824 if node.dyn_kwargs: 

825 self.write(", **") 

826 self.visit(node.dyn_kwargs, frame) 

827 

828 def visit_Assign(self, node: nodes.Assign, frame: Frame) -> None: 

829 """Visit an assignment node ({% set %} statements).""" 

830 frame = t.cast(AsyncFrame, frame) 

831 self.push_assign_tracking() 

832 

833 # Check for namespace assignments like `ns.var = value` 

834 seen_refs: set[str] = set() 

835 for nsref in node.find_all(nodes.NSRef): 

836 if nsref.name in seen_refs: 

837 continue 

838 seen_refs.add(nsref.name) 

839 ref = frame.symbols.ref(nsref.name) 

840 self.writeline(f"if not isinstance({ref}, Namespace):") 

841 self.indent() 

842 self.writeline( 

843 "raise TemplateRuntimeError" 

844 '("cannot assign attribute on non-namespace object")' 

845 ) 

846 self.outdent() 

847 

848 # Generate the assignment code 

849 self.newline(node) 

850 self.visit(node.target, frame) 

851 self.write(" = ") 

852 self.visit(node.node, frame) 

853 self.pop_assign_tracking(frame) 

854 

855 def push_assign_tracking(self) -> None: 

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

857 self._assign_stack.append(set()) 

858 

859 def pop_assign_tracking(self, frame: Frame) -> None: 

860 """Pop the topmost level for assignment tracking and update context variables.""" 

861 frame = t.cast(AsyncFrame, frame) 

862 vars_set = self._assign_stack.pop() 

863 

864 if ( 

865 not frame.block_frame 

866 and not frame.loop_frame 

867 and not frame.toplevel 

868 or not vars_set 

869 ): 

870 return 

871 

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

873 

874 if len(vars_set) == 1: 

875 name = next(iter(vars_set)) 

876 ref = frame.symbols.ref(name) 

877 if frame.loop_frame: 

878 self.writeline(f"_loop_vars[{name!r}] = {ref}") 

879 return 

880 if frame.block_frame: 

881 self.writeline(f"_block_vars[{name!r}] = {ref}") 

882 return 

883 self.writeline(f"context.vars[{name!r}] = {ref}") 

884 else: 

885 if frame.loop_frame: 

886 self.writeline("_loop_vars.update({") 

887 elif frame.block_frame: 

888 self.writeline("_block_vars.update({") 

889 else: 

890 self.writeline("context.vars.update({") 

891 for idx, name in enumerate(sorted(vars_set)): 

892 if idx: 

893 self.write(", ") 

894 ref = frame.symbols.ref(name) 

895 self.write(f"{name!r}: {ref}") 

896 self.write("})") 

897 

898 if not frame.block_frame and not frame.loop_frame and public_names: 

899 if len(public_names) == 1: 

900 self.writeline(f"context.exported_vars.add({public_names[0]!r})") 

901 else: 

902 names_str = ", ".join(map(repr, sorted(public_names))) 

903 self.writeline(f"context.exported_vars.update(({names_str}))")