variables.py 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. """
  2. Core variables support for interactive collection and detection.
  3. Provides a BaseVariables class that can detect which variable sets are used
  4. in a Jinja2 template and interactively collect values from the user.
  5. """
  6. from typing import Dict, List, Tuple, Set, Any
  7. import jinja2
  8. from jinja2 import meta
  9. import typer
  10. from .prompt import PromptHandler
  11. class BaseVariables:
  12. """Base implementation for variable sets and interactive prompting.
  13. Subclasses should set `variable_sets` to one of two shapes:
  14. 1) Legacy shape (mapping of set-name -> { var_name: { ... } })
  15. { "general": { "foo": { ... }, ... } }
  16. 2) New shape (mapping of set-name -> { "prompt": str, "variables": { var_name: { ... } } })
  17. { "general": { "prompt": "...", "variables": { "foo": { ... } } } }
  18. """
  19. variable_sets: Dict[str, Dict[str, Any]] = {}
  20. def __init__(self) -> None:
  21. # Flattened list of all declared variable names -> (set_name, meta)
  22. self._declared: Dict[str, Tuple[str, Dict[str, Any]]] = {}
  23. # Support both legacy and new shapes. If the set value contains a
  24. # 'variables' key, use that mapping; otherwise assume the mapping is
  25. # directly the vars map (legacy).
  26. for set_name, set_def in getattr(self, "variable_sets", {}).items():
  27. vars_map = set_def.get("variables") if isinstance(set_def, dict) and "variables" in set_def else set_def
  28. if not isinstance(vars_map, dict):
  29. continue
  30. for var_name, meta_info in vars_map.items():
  31. self._declared[var_name] = (set_name, meta_info)
  32. def find_used_variables(self, template_content: str) -> Set[str]:
  33. """Parse the Jinja2 template and return the set of variable names used."""
  34. env = jinja2.Environment()
  35. try:
  36. ast = env.parse(template_content)
  37. used = meta.find_undeclared_variables(ast)
  38. return set(used)
  39. except Exception:
  40. # If parsing fails, fallback to an empty set (safe behavior)
  41. return set()
  42. def find_used_subscript_keys(self, template_content: str) -> Dict[str, Set[str]]:
  43. """Return mapping of variable name -> set of string keys accessed via subscripting
  44. Example: for template using service_port['http'] and service_port['https']
  45. this returns { 'service_port': {'http', 'https'} }.
  46. """
  47. try:
  48. env = jinja2.Environment()
  49. ast = env.parse(template_content)
  50. # Walk AST and collect Subscript nodes
  51. from jinja2 import nodes
  52. subs: Dict[str, Set[str]] = {}
  53. for node in ast.find_all(nodes.Getitem):
  54. # Getitem node structure: node.node (value), node.arg (index)
  55. try:
  56. if isinstance(node.node, nodes.Name):
  57. var_name = node.node.name
  58. # index can be Const (string) or Name/other; handle Const
  59. idx = node.arg
  60. if isinstance(idx, nodes.Const) and isinstance(idx.value, str):
  61. subs.setdefault(var_name, set()).add(idx.value)
  62. except Exception:
  63. continue
  64. return subs
  65. except Exception:
  66. return {}
  67. def extract_template_defaults(self, template_content: str) -> Dict[str, Any]:
  68. """Extract default values from Jinja2 expressions like {{ var | default(value) }}."""
  69. import re
  70. def _parse_literal(s: str):
  71. s = s.strip()
  72. if s.startswith("'") and s.endswith("'"):
  73. return s[1:-1]
  74. if s.startswith('"') and s.endswith('"'):
  75. return s[1:-1]
  76. if s.isdigit():
  77. return int(s)
  78. return s
  79. defaults: Dict[str, Any] = {}
  80. # Match {{ var['key'] | default(value) }} and {{ var | default(value) }}
  81. pattern_subscript = r'\{\{\s*(\w+)\s*\[\s*["\']([^"\']+)["\']\s*\]\s*\|\s*default\(([^)]+)\)\s*\}\}'
  82. for var, key, default_str in re.findall(pattern_subscript, template_content):
  83. if var not in defaults or not isinstance(defaults[var], dict):
  84. defaults[var] = {}
  85. defaults[var][key] = _parse_literal(default_str)
  86. pattern_scalar = r'\{\{\s*(\w+)\s*\|\s*default\(([^)]+)\)\s*\}\}'
  87. for var, default_str in re.findall(pattern_scalar, template_content):
  88. # Only set scalar default if not already set as a dict
  89. if var not in defaults:
  90. defaults[var] = _parse_literal(default_str)
  91. # Handle simple {% set name = other | default('val') %} patterns
  92. set_pattern = r"\{%\s*set\s+(\w+)\s*=\s*([^%]+?)\s*%}"
  93. for set_var, expr in re.findall(set_pattern, template_content):
  94. m = re.match(r"(\w+)\s*\|\s*default\(([^)]+)\)", expr.strip())
  95. if m:
  96. src_var, src_default = m.groups()
  97. if src_var in defaults:
  98. defaults[set_var] = defaults[src_var]
  99. else:
  100. defaults[set_var] = _parse_literal(src_default)
  101. # Resolve transitive references: if a default is an identifier that
  102. # points to another default, follow it; if it points to a declared
  103. # variable with a metadata default, use that.
  104. def _resolve_ref(value, seen: Set[str]):
  105. if not isinstance(value, str):
  106. return value
  107. if value in seen:
  108. return value
  109. seen.add(value)
  110. if value in defaults:
  111. return _resolve_ref(defaults[value], seen)
  112. if value in self._declared:
  113. declared_def = self._declared[value][1].get("default")
  114. if declared_def is not None:
  115. return declared_def
  116. return value
  117. for k in list(defaults.keys()):
  118. defaults[k] = _resolve_ref(defaults[k], set([k]))
  119. return defaults
  120. def extract_variable_meta_overrides(self, template_content: str) -> Dict[str, Dict[str, Any]]:
  121. """Extract variable metadata overrides from a Jinja2 block.
  122. Supports a block like:
  123. {% variables %}
  124. container_hostname:
  125. description: "..."
  126. {% endvariables %}
  127. The contents are parsed as YAML and returned as a dict mapping
  128. variable name -> metadata overrides.
  129. """
  130. import re
  131. try:
  132. m = re.search(r"\{%\s*variables\s*%\}(.+?)\{%\s*endvariables\s*%\}", template_content, flags=re.S)
  133. if not m:
  134. return {}
  135. yaml_block = m.group(1).strip()
  136. try:
  137. import yaml
  138. except Exception:
  139. return {}
  140. try:
  141. data = yaml.safe_load(yaml_block) or {}
  142. if isinstance(data, dict):
  143. # Ensure values are dicts
  144. cleaned: Dict[str, Dict[str, Any]] = {}
  145. for k, v in data.items():
  146. if v is None:
  147. cleaned[k] = {}
  148. elif isinstance(v, dict):
  149. cleaned[k] = v
  150. else:
  151. # If a scalar was provided, interpret as description
  152. cleaned[k] = {"description": v}
  153. return cleaned
  154. except Exception:
  155. return {}
  156. except Exception:
  157. return {}
  158. return {}
  159. def determine_variable_sets(self, template_content: str) -> Tuple[List[str], Set[str]]:
  160. """Return a list of variable set names that contain any used variables.
  161. Also returns the raw set of used variable names.
  162. """
  163. used = self.find_used_variables(template_content)
  164. matched_sets: List[str] = []
  165. for set_name, set_def in getattr(self, "variable_sets", {}).items():
  166. vars_map = set_def.get("variables") if isinstance(set_def, dict) and "variables" in set_def else set_def
  167. if not isinstance(vars_map, dict):
  168. continue
  169. if any(var in used for var in vars_map.keys()):
  170. matched_sets.append(set_name)
  171. return matched_sets, used
  172. def collect_values(self, used_vars: Set[str], template_defaults: Dict[str, Any] = None, used_subscripts: Dict[str, Set[str]] = None) -> Dict[str, Any]:
  173. """Interactively prompt for values for the variables that appear in the template.
  174. For variables that were declared in `variable_sets` we use their metadata.
  175. For unknown variables, we fall back to a generic prompt.
  176. """
  177. prompt_handler = PromptHandler(self._declared, getattr(self, "variable_sets", {}))
  178. return prompt_handler.collect_values(used_vars, template_defaults, used_subscripts)