Coverage for promplate/prompt/builder.py: 90%

29 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-08-08 04:34 +0800

1from types import FunctionType 

2 

3 

4class CodeBuilder: 

5 """Build source code conveniently.""" 

6 

7 def __init__(self, indent_level=0, indent_str="\t"): 

8 self.code = [] 

9 self.indent_level = indent_level 

10 self.indent_str = indent_str 

11 

12 def __str__(self): 

13 return "".join(map(str, self.code)) 

14 

15 def add_line(self, line): 

16 """Add a line of source to the code.""" 

17 self.code.extend((self.indent_str * self.indent_level, line, "\n")) 

18 

19 return self 

20 

21 def add_section(self): 

22 """Add a section, a sub-CodeBuilder.""" 

23 section = CodeBuilder(self.indent_level, self.indent_str) 

24 self.code.append(section) 

25 

26 return section 

27 

28 def indent(self): 

29 """Increase the current indent for following lines.""" 

30 self.indent_level += 1 

31 

32 return self 

33 

34 def dedent(self): 

35 """Decrease the current indent for following lines.""" 

36 self.indent_level -= 1 

37 

38 return self 

39 

40 def get_render_function(self) -> FunctionType: 

41 """Execute the code, and return a dict of globals it defines.""" 

42 assert self.indent_level == 0 

43 python_source = str(self) 

44 global_namespace = {} 

45 exec(python_source, global_namespace) 

46 return global_namespace["render"] 

47 

48 

49def get_base_builder(sync=True, indent_str="\t"): 

50 return ( 

51 CodeBuilder(indent_str=indent_str) 

52 .add_line("def render():" if sync else "async def render():") 

53 .indent() 

54 .add_line("__parts__ = []") 

55 .add_line("__append__ = __parts__.append") 

56 )