Coverage for jinja2_async_environment/compiler_modules/loops.py: 85%

1"""Loop code generation utilities for for-loops.""" 

2 

3import typing as t 

4 

5from jinja2 import nodes 

6 

7if t.TYPE_CHECKING: 

8 from .codegen import AsyncCodeGenerator 

9 from .frame import AsyncFrame 

10 

11 

12class LoopCodeGenerator: 

13 """Utility class for generating for-loop specific code patterns.""" 

14 

15 def __init__(self, code_generator: "AsyncCodeGenerator") -> None: 

16 self.code_generator = code_generator 

17 

18 def generate_async_for_header( 

19 self, node: nodes.For, target: nodes.Node, frame: "AsyncFrame" 

20 ) -> None: 

21 """Generate the async for loop header with proper syntax.""" 

22 self.code_generator.writeline( 

23 self.code_generator.choose_async("async for ", "for "), node 

24 ) 

25 self.code_generator.visit(target, frame) 

26 

27 def generate_loop_iterator( 

28 self, 

29 iter_node: nodes.Node, 

30 frame: "AsyncFrame", 

31 extended_loop: bool, 

32 loop_ref: str | None, 

33 loop_filter_func: str | None, 

34 ) -> None: 

35 """Generate the iterator part of the for loop.""" 

36 if extended_loop and loop_ref: 

37 self.code_generator.write( 

38 f", {loop_ref} in {self.code_generator.choose_async('Async')}LoopContext(" 

39 ) 

40 else: 

41 self.code_generator.write(" in ") 

42 

43 if loop_filter_func: 

44 self.code_generator.write(f"{loop_filter_func}(") 

45 

46 if self.code_generator.environment.is_async and not extended_loop: 

47 self.code_generator.write("auto_aiter(") 

48 

49 self.code_generator.visit(iter_node, frame) 

50 

51 if self.code_generator.environment.is_async and not extended_loop: 

52 self.code_generator.write(")") 

53 

54 if loop_filter_func: 

55 self.code_generator.write(")") 

56 

57 if extended_loop: 

58 self.code_generator.write(", undefined):") 

59 else: 

60 self.code_generator.write(":")