Coverage for /usr/local/lib/python3.12/site-packages/prefect/utilities/templating.py: 15%
164 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-12-05 13:38 +0000
1import enum 1a
2import os 1a
3import re 1a
4from typing import ( 1a
5 TYPE_CHECKING,
6 Any,
7 Callable,
8 Literal,
9 NamedTuple,
10 Optional,
11 TypeVar,
12 Union,
13 cast,
14 overload,
15)
17from prefect.client.utilities import inject_client 1a
18from prefect.logging.loggers import get_logger 1a
19from prefect.utilities.annotations import NotSet 1a
20from prefect.utilities.collections import get_from_dict 1a
22if TYPE_CHECKING: 22 ↛ 23line 22 didn't jump to line 23 because the condition on line 22 was never true1a
23 import logging
25 from prefect.client.orchestration import PrefectClient
28T = TypeVar("T", str, int, float, bool, dict[Any, Any], list[Any], None) 1a
30PLACEHOLDER_CAPTURE_REGEX = re.compile(r"({{\s*([\w\.\-\[\]$]+)\s*}})") 1a
31BLOCK_DOCUMENT_PLACEHOLDER_PREFIX = "prefect.blocks." 1a
32VARIABLE_PLACEHOLDER_PREFIX = "prefect.variables." 1a
33ENV_VAR_PLACEHOLDER_PREFIX = "$" 1a
36logger: "logging.Logger" = get_logger("utilities.templating") 1a
39class PlaceholderType(enum.Enum): 1a
40 STANDARD = "standard" 1a
41 BLOCK_DOCUMENT = "block_document" 1a
42 VARIABLE = "variable" 1a
43 ENV_VAR = "env_var" 1a
46class Placeholder(NamedTuple): 1a
47 full_match: str 1a
48 name: str 1a
49 type: PlaceholderType 1a
52def determine_placeholder_type(name: str) -> PlaceholderType: 1a
53 """
54 Determines the type of a placeholder based on its name.
56 Args:
57 name: The name of the placeholder
59 Returns:
60 The type of the placeholder
61 """
62 if name.startswith(BLOCK_DOCUMENT_PLACEHOLDER_PREFIX):
63 return PlaceholderType.BLOCK_DOCUMENT
64 elif name.startswith(VARIABLE_PLACEHOLDER_PREFIX):
65 return PlaceholderType.VARIABLE
66 elif name.startswith(ENV_VAR_PLACEHOLDER_PREFIX):
67 return PlaceholderType.ENV_VAR
68 else:
69 return PlaceholderType.STANDARD
72def find_placeholders(template: T) -> set[Placeholder]: 1a
73 """
74 Finds all placeholders in a template.
76 Args:
77 template: template to discover placeholders in
79 Returns:
80 A set of all placeholders in the template
81 """
82 seed: set[Placeholder] = set()
83 if isinstance(template, (int, float, bool)):
84 return seed
85 if isinstance(template, str):
86 result = PLACEHOLDER_CAPTURE_REGEX.findall(template)
87 return {
88 Placeholder(full_match, name, determine_placeholder_type(name))
89 for full_match, name in result
90 }
91 elif isinstance(template, dict):
92 return seed.union(*[find_placeholders(value) for value in template.values()])
93 elif isinstance(template, list):
94 return seed.union(*[find_placeholders(item) for item in template])
95 else:
96 raise ValueError(f"Unexpected type: {type(template)}")
99@overload 1a
100def apply_values( 100 ↛ exitline 100 didn't return from function 'apply_values' because 1a
101 template: T,
102 values: dict[str, Any],
103 remove_notset: Literal[True] = True,
104 warn_on_notset: bool = False,
105) -> T: ...
108@overload 1a
109def apply_values( 109 ↛ exitline 109 didn't return from function 'apply_values' because 1a
110 template: T,
111 values: dict[str, Any],
112 remove_notset: Literal[False] = False,
113 warn_on_notset: bool = False,
114) -> Union[T, type[NotSet]]: ...
117@overload 1a
118def apply_values( 118 ↛ exitline 118 didn't return from function 'apply_values' because 1a
119 template: T,
120 values: dict[str, Any],
121 remove_notset: bool = False,
122 warn_on_notset: bool = False,
123) -> Union[T, type[NotSet]]: ...
126def apply_values( 1a
127 template: T,
128 values: dict[str, Any],
129 remove_notset: bool = True,
130 warn_on_notset: bool = False,
131) -> Union[T, type[NotSet]]:
132 """
133 Replaces placeholders in a template with values from a supplied dictionary.
135 Will recursively replace placeholders in dictionaries and lists.
137 If a value has no placeholders, it will be returned unchanged.
139 If a template contains only a single placeholder, the placeholder will be
140 fully replaced with the value.
142 If a template contains text before or after a placeholder or there are
143 multiple placeholders, the placeholders will be replaced with the
144 corresponding variable values.
146 If a template contains a placeholder that is not in `values`, NotSet will
147 be returned to signify that no placeholder replacement occurred. If
148 `template` is a dictionary that contains a key with a value of NotSet,
149 the key will be removed in the return value unless `remove_notset` is set to False.
151 Args:
152 template: template to discover and replace values in
153 values: The values to apply to placeholders in the template
154 remove_notset: If True, remove keys with an unset value
155 warn_on_notset: If True, warn when a placeholder is not found in `values`
157 Returns:
158 The template with the values applied
159 """
160 if template in (NotSet, None) or isinstance(template, (int, float)):
161 return template
162 if isinstance(template, str):
163 placeholders = find_placeholders(template)
164 if not placeholders:
165 # If there are no values, we can just use the template
166 return template
167 elif (
168 len(placeholders) == 1
169 and list(placeholders)[0].full_match == template
170 and list(placeholders)[0].type is PlaceholderType.STANDARD
171 ):
172 # If there is only one variable with no surrounding text,
173 # we can replace it. If there is no variable value, we
174 # return NotSet to indicate that the value should not be included.
175 value = get_from_dict(values, list(placeholders)[0].name, NotSet)
176 if value is NotSet and warn_on_notset:
177 logger.warning(
178 f"Value for placeholder {list(placeholders)[0].name!r} not found in provided values. Please ensure that "
179 "the placeholder is spelled correctly and that the corresponding value is provided.",
180 )
181 return value
182 else:
183 for full_match, name, placeholder_type in placeholders:
184 if placeholder_type is PlaceholderType.STANDARD:
185 value = get_from_dict(values, name, NotSet)
186 elif placeholder_type is PlaceholderType.ENV_VAR:
187 name = name.lstrip(ENV_VAR_PLACEHOLDER_PREFIX)
188 value = os.environ.get(name, NotSet)
189 else:
190 continue
192 if value is NotSet:
193 if warn_on_notset:
194 logger.warning(
195 f"Value for placeholder {full_match!r} not found in provided values. Please ensure that "
196 "the placeholder is spelled correctly and that the corresponding value is provided.",
197 )
198 if remove_notset:
199 template = template.replace(full_match, "")
200 else:
201 template = template.replace(full_match, str(value))
203 return template
204 elif isinstance(template, dict):
205 updated_template: dict[str, Any] = {}
206 for key, value in template.items():
207 updated_value = apply_values(
208 value,
209 values,
210 remove_notset=remove_notset,
211 warn_on_notset=warn_on_notset,
212 )
213 if updated_value is not NotSet:
214 updated_template[key] = updated_value
215 elif not remove_notset:
216 updated_template[key] = value
218 return cast(T, updated_template)
219 elif isinstance(template, list):
220 updated_list: list[Any] = []
221 for value in template:
222 updated_value = apply_values(
223 value,
224 values,
225 remove_notset=remove_notset,
226 warn_on_notset=warn_on_notset,
227 )
228 if updated_value is not NotSet:
229 updated_list.append(updated_value)
230 return cast(T, updated_list)
231 else:
232 raise ValueError(f"Unexpected template type {type(template).__name__!r}")
235@inject_client 1a
236async def resolve_block_document_references( 1a
237 template: T,
238 client: Optional["PrefectClient"] = None,
239 value_transformer: Optional[Callable[[str, Any], Any]] = None,
240) -> Union[T, dict[str, Any]]:
241 """
242 Resolve block document references in a template by replacing each reference with
243 its value or the return value of the transformer function if provided.
245 Recursively searches for block document references in dictionaries and lists.
247 Identifies block document references by the as dictionary with the following
248 structure:
249 ```
250 {
251 "$ref": {
252 "block_document_id": <block_document_id>
253 }
254 }
255 ```
256 where `<block_document_id>` is the ID of the block document to resolve.
258 Once the block document is retrieved from the API, the data of the block document
259 is used to replace the reference.
261 Accessing Values:
262 -----------------
263 To access different values in a block document, use dot notation combined with the block document's prefix, slug, and block name.
265 For a block document with the structure:
266 ```json
267 {
268 "value": {
269 "key": {
270 "nested-key": "nested-value"
271 },
272 "list": [
273 {"list-key": "list-value"},
274 1,
275 2
276 ]
277 }
278 }
279 ```
280 examples of value resolution are as follows:
282 1. Accessing a nested dictionary:
283 Format: `prefect.blocks.<block_type_slug>.<block_document_name>.value.key`
284 Example: Returns `{"nested-key": "nested-value"}`
286 2. Accessing a specific nested value:
287 Format: `prefect.blocks.<block_type_slug>.<block_document_name>.value.key.nested-key`
288 Example: Returns `"nested-value"`
290 3. Accessing a list element's key-value:
291 Format: `prefect.blocks.<block_type_slug>.<block_document_name>.value.list[0].list-key`
292 Example: Returns `"list-value"`
294 Default Resolution for System Blocks:
295 -------------------------------------
296 For system blocks, which only contain a `value` attribute, this attribute is resolved by default.
298 Args:
299 template: The template to resolve block documents in
300 value_transformer: A function that takes the block placeholder and the block value and returns replacement text for the template
302 Returns:
303 The template with block documents resolved
304 """
305 if TYPE_CHECKING:
306 # The @inject_client decorator takes care of providing the client, but
307 # the function signature must mark it as optional to callers.
308 assert client is not None
310 if isinstance(template, dict):
311 block_document_id = template.get("$ref", {}).get("block_document_id")
312 if block_document_id:
313 block_document = await client.read_block_document(block_document_id)
314 return block_document.data
315 updated_template: dict[str, Any] = {}
316 for key, value in template.items():
317 updated_value = await resolve_block_document_references(
318 value, value_transformer=value_transformer, client=client
319 )
320 updated_template[key] = updated_value
321 return updated_template
322 elif isinstance(template, list):
323 return [
324 await resolve_block_document_references(
325 item, value_transformer=value_transformer, client=client
326 )
327 for item in template
328 ]
329 elif isinstance(template, str):
330 placeholders = find_placeholders(template)
331 has_block_document_placeholder = any(
332 placeholder.type is PlaceholderType.BLOCK_DOCUMENT
333 for placeholder in placeholders
334 )
335 if not (placeholders and has_block_document_placeholder):
336 return template
337 elif (
338 len(placeholders) == 1
339 and list(placeholders)[0].full_match == template
340 and list(placeholders)[0].type is PlaceholderType.BLOCK_DOCUMENT
341 ):
342 # value_keypath will be a list containing a dot path if additional
343 # attributes are accessed and an empty list otherwise.
344 [placeholder] = placeholders
345 parts = placeholder.name.replace(
346 BLOCK_DOCUMENT_PLACEHOLDER_PREFIX, ""
347 ).split(".", 2)
348 block_type_slug, block_document_name, *value_keypath = parts
349 block_document = await client.read_block_document_by_name(
350 name=block_document_name, block_type_slug=block_type_slug
351 )
352 data = block_document.data
353 value: Union[T, dict[str, Any]] = data
355 # resolving system blocks to their data for backwards compatibility
356 if len(data) == 1 and "value" in data:
357 # only resolve the value if the keypath is not already pointing to "value"
358 if not (value_keypath and value_keypath[0].startswith("value")):
359 data = value = value["value"]
361 # resolving keypath/block attributes
362 if value_keypath:
363 from_dict: Any = get_from_dict(data, value_keypath[0], default=NotSet)
364 if from_dict is NotSet:
365 raise ValueError(
366 f"Invalid template: {template!r}. Could not resolve the"
367 " keypath in the block document data."
368 )
369 value = from_dict
371 if value_transformer:
372 value = value_transformer(placeholder.full_match, value)
374 return value
375 else:
376 raise ValueError(
377 f"Invalid template: {template!r}. Only a single block placeholder is"
378 " allowed in a string and no surrounding text is allowed."
379 )
381 return template
384@inject_client 1a
385async def resolve_variables(template: T, client: Optional["PrefectClient"] = None) -> T: 1a
386 """
387 Resolve variables in a template by replacing each variable placeholder with the
388 value of the variable.
390 Recursively searches for variable placeholders in dictionaries and lists.
392 Strips variable placeholders if the variable is not found.
394 Args:
395 template: The template to resolve variables in
397 Returns:
398 The template with variables resolved
399 """
400 if TYPE_CHECKING:
401 # The @inject_client decorator takes care of providing the client, but
402 # the function signature must mark it as optional to callers.
403 assert client is not None
405 if isinstance(template, str):
406 placeholders = find_placeholders(template)
407 has_variable_placeholder = any(
408 placeholder.type is PlaceholderType.VARIABLE for placeholder in placeholders
409 )
410 if not placeholders or not has_variable_placeholder:
411 # If there are no values, we can just use the template
412 return template
413 elif (
414 len(placeholders) == 1
415 and list(placeholders)[0].full_match == template
416 and list(placeholders)[0].type is PlaceholderType.VARIABLE
417 ):
418 variable_name = list(placeholders)[0].name.replace(
419 VARIABLE_PLACEHOLDER_PREFIX, ""
420 )
421 variable = await client.read_variable_by_name(name=variable_name)
422 if variable is None:
423 return ""
424 else:
425 return cast(T, variable.value)
426 else:
427 for full_match, name, placeholder_type in placeholders:
428 if placeholder_type is PlaceholderType.VARIABLE:
429 variable_name = name.replace(VARIABLE_PLACEHOLDER_PREFIX, "")
430 variable = await client.read_variable_by_name(name=variable_name)
431 if variable is None:
432 template = template.replace(full_match, "")
433 else:
434 template = template.replace(full_match, str(variable.value))
435 return template
436 elif isinstance(template, dict):
437 return {
438 key: await resolve_variables(value, client=client)
439 for key, value in template.items()
440 }
441 elif isinstance(template, list):
442 return [await resolve_variables(item, client=client) for item in template]
443 else:
444 return template