Coverage for src/kwai/core/template/jinja2_engine.py: 78%

1"""Modules that implements the template engine interface for jinja2.""" 

2 

3from typing import Any 

4 

5from fastapi import Request 

6from fastapi.templating import Jinja2Templates 

7from jinja2 import ( 

8 ChoiceLoader, 

9 Environment, 

10 PackageLoader, 

11 PrefixLoader, 

12 TemplatesNotFound, 

13 select_autoescape, 

14) 

15 

16from .jinja2_template import Jinja2Template 

17from .template import Template 

18from .template_engine import TemplateEngine, TemplateNotFoundException 

19 

20 

21JINJA2_FILE_EXTENSION = ".jinja2" 

22 

23 

24class Jinja2Engine(TemplateEngine): 

25 """Implements the TemplateEngine interface for Jinja2.""" 

26 

27 def __init__(self, **kwargs): 

28 """Construct a Jinja2Engine. 

29 

30 kwargs will be merged to the variables that are used to render a template. 

31 Use it for variables that are used in all templates. 

32 """ 

33 self._env = Environment( 

34 loader=ChoiceLoader( 

35 [ 

36 PrefixLoader( 

37 { 

38 "nl": PackageLoader("kwai", "templates"), 

39 "en": PackageLoader("kwai", "templates"), 

40 } 

41 ), 

42 PackageLoader("kwai", "templates/nl"), 

43 PackageLoader("kwai", "templates"), 

44 ], 

45 ), 

46 autoescape=select_autoescape( 

47 disabled_extensions=("txt",), 

48 default_for_string=True, 

49 default=True, 

50 ), 

51 extensions=["jinja2.ext.do"], 

52 ) 

53 self._variables = kwargs 

54 

55 @property 

56 def web_templates(self) -> Jinja2Templates: 

57 """Return templates that can be used with a TemplateResponse.""" 

58 

59 def app_context(request: Request) -> dict[str, Any]: 

60 return self._variables 

61 

62 return Jinja2Templates(env=self._env, context_processors=[app_context]) 

63 

64 def create(self, template_file_path: str) -> Template: 

65 """Create a jinja2 template.""" 

66 try: 

67 template = self._env.get_template(template_file_path + ".jinja2") 

68 except TemplatesNotFound as exc: 

69 raise TemplateNotFoundException( 

70 f"Could not find a template with name '{template_file_path}'" 

71 ) from exc 

72 

73 return Jinja2Template(template, **self._variables)