Coverage for haystack/utils/jinja2_extensions.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any 

6 

7from jinja2 import Environment, meta, nodes 

8from jinja2.ext import Extension 

9 

10from haystack.lazy_imports import LazyImport 

11 

12with LazyImport(message='Run "pip install arrow>=1.3.0"') as arrow_import: 

13 import arrow 

14 

15 

16class Jinja2TimeExtension(Extension): 

17 """A Jinja2 extension for formatting dates and times.""" 

18 

19 # Syntax for current date 

20 tags = {"now"} 

21 

22 def __init__(self, environment: Environment) -> None: 

23 """ 

24 Initializes the JinjaTimeExtension object. 

25 

26 :param environment: The Jinja2 environment to initialize the extension with. 

27 It provides the context where the extension will operate. 

28 """ 

29 arrow_import.check() 

30 super().__init__(environment) 

31 

32 @staticmethod 

33 def _get_datetime( 

34 timezone: str, operator: str | None = None, offset: str | None = None, datetime_format: str | None = None 

35 ) -> str: 

36 """ 

37 Get the current datetime based on timezone, apply any offset if provided, and format the result. 

38 

39 :param timezone: The timezone string (e.g., 'UTC' or 'America/New_York') for which the current 

40 time should be fetched. 

41 :param operator: The operator ('+' or '-') to apply to the offset (used for adding/subtracting intervals). 

42 Defaults to None if no offset is applied, otherwise default is '+'. 

43 :param offset: The offset string in the format 'interval=value' (e.g., 'hours=2,days=1') specifying how much 

44 to adjust the datetime. The intervals can be any valid interval accepted 

45 by Arrow (e.g., hours, days, weeks, months). Defaults to None if no adjustment is needed. 

46 :param datetime_format: The format string to use for formatting the output datetime. 

47 Defaults to '%Y-%m-%d %H:%M:%S' if not provided. 

48 """ 

49 try: 

50 dt = arrow.now(timezone) 

51 except Exception as e: 

52 raise ValueError(f"Invalid timezone {timezone}: {e}") from e 

53 

54 if offset and operator: 

55 try: 

56 # Parse the offset and apply it to the datetime object 

57 replace_params: dict[str, Any] = { 

58 interval.strip(): float(operator + value.strip()) 

59 for param in offset.split(",") 

60 for interval, value in [param.split("=")] 

61 } 

62 # Shift the datetime fields based on the parsed offset 

63 dt = dt.shift(**replace_params) 

64 except (ValueError, AttributeError) as e: 

65 raise ValueError(f"Invalid offset or operator {offset}, {operator}: {e}") from e 

66 

67 # Use the provided format or fallback to the default one 

68 datetime_format = datetime_format or "%Y-%m-%d %H:%M:%S" 

69 

70 return dt.strftime(datetime_format) 

71 

72 def parse(self, parser: Any) -> nodes.Node | list[nodes.Node]: 

73 """ 

74 Parse the template expression to determine how to handle the datetime formatting. 

75 

76 :param parser: The parser object that processes the template expressions and manages the syntax tree. 

77 It's used to interpret the template's structure. 

78 """ 

79 lineno = next(parser.stream).lineno 

80 node = parser.parse_expression() 

81 # Check if a custom datetime format is provided after a comma 

82 datetime_format = parser.parse_expression() if parser.stream.skip_if("comma") else nodes.Const(None) 

83 

84 # Default Add when no operator is provided 

85 operator = "+" if isinstance(node, nodes.Add) else "-" 

86 # Call the _get_datetime method with the appropriate operator and offset, if exist 

87 call_method = self.call_method( 

88 "_get_datetime", 

89 [node.left, nodes.Const(operator), node.right, datetime_format] 

90 if isinstance(node, (nodes.Add, nodes.Sub)) 

91 else [node, nodes.Const(None), nodes.Const(None), datetime_format], 

92 lineno=lineno, 

93 ) 

94 

95 return nodes.Output([call_method], lineno=lineno) 

96 

97 

98def _collect_assigned_variables(ast: nodes.Template) -> set[str]: 

99 """ 

100 Extract variables assigned within the Jinja2 template AST. 

101 

102 :param ast: The Jinja2 Abstract Syntax Tree (AST) of the template. 

103 

104 :returns: 

105 A set of variable names that are assigned within the template. 

106 """ 

107 # Collect all variables assigned inside the template via {% set %} 

108 assigned_variables = set() 

109 

110 for node in ast.find_all(nodes.Assign): 

111 if isinstance(node.target, nodes.Name): 

112 assigned_variables.add(node.target.name) 

113 elif isinstance(node.target, (nodes.List, nodes.Tuple)): 

114 for name_node in node.target.items: 

115 if isinstance(name_node, nodes.Name): 

116 assigned_variables.add(name_node.name) 

117 

118 return assigned_variables 

119 

120 

121def _extract_template_variables_and_assignments(env: Environment, template: str) -> tuple[set[str], set[str]]: 

122 """ 

123 Extract variables from a Jinja2 template and variables assigned within it. 

124 

125 :param env: A Jinja2 environment. 

126 :param template: A Jinja2 template string. 

127 :returns: A tuple of (assigned_variables, template_variables) where: 

128 - assigned_variables: Variables assigned within the template (e.g., via {% set %}) 

129 - template_variables: All undeclared variables used in the template 

130 """ 

131 jinja2_ast = env.parse(template) 

132 template_variables = meta.find_undeclared_variables(jinja2_ast) 

133 assigned_variables = _collect_assigned_variables(jinja2_ast) 

134 return assigned_variables, template_variables