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

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

2

3 import typing as t

4

5 from jinja2 import nodes

6

7 if t.TYPE_CHECKING:

8 from .codegen import AsyncCodeGenerator

9 from .frame import AsyncFrame

10

11

12 class 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(":")