prefect-client 3.1.5__py3-none-any.whl → 3.1.7__py3-none-any.whl

prefect/utilities/templating.py

@@ -1,17 +1,7 @@
 import enum
 import os
 import re
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    NamedTuple,
-    Optional,
-    Set,
-    Type,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, NamedTuple, Optional, TypeVar, Union, cast
 
 from prefect.client.utilities import inject_client
 from prefect.utilities.annotations import NotSet
@@ -21,7 +11,7 @@ if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
 
 
-T = TypeVar("T", str, int, float, bool, dict, list, None)
+T = TypeVar("T", str, int, float, bool, dict[Any, Any], list[Any], None)
 
 PLACEHOLDER_CAPTURE_REGEX = re.compile(r"({{\s*([\w\.\-\[\]$]+)\s*}})")
 BLOCK_DOCUMENT_PLACEHOLDER_PREFIX = "prefect.blocks."
@@ -62,7 +52,7 @@ def determine_placeholder_type(name: str) -> PlaceholderType:
         return PlaceholderType.STANDARD
 
 
-def find_placeholders(template: T) -> Set[Placeholder]:
+def find_placeholders(template: T) -> set[Placeholder]:
     """
     Finds all placeholders in a template.
 
@@ -72,8 +62,9 @@ def find_placeholders(template: T) -> Set[Placeholder]:
     Returns:
         A set of all placeholders in the template
     """
+    seed: set[Placeholder] = set()
     if isinstance(template, (int, float, bool)):
-        return set()
+        return seed
     if isinstance(template, str):
         result = PLACEHOLDER_CAPTURE_REGEX.findall(template)
         return {
@@ -81,18 +72,16 @@ def find_placeholders(template: T) -> Set[Placeholder]:
             for full_match, name in result
         }
     elif isinstance(template, dict):
-        return set().union(
-            *[find_placeholders(value) for key, value in template.items()]
-        )
+        return seed.union(*[find_placeholders(value) for value in template.values()])
     elif isinstance(template, list):
-        return set().union(*[find_placeholders(item) for item in template])
+        return seed.union(*[find_placeholders(item) for item in template])
     else:
         raise ValueError(f"Unexpected type: {type(template)}")
 
 
 def apply_values(
-    template: T, values: Dict[str, Any], remove_notset: bool = True
-) -> Union[T, Type[NotSet]]:
+    template: T, values: dict[str, Any], remove_notset: bool = True
+) -> Union[T, type[NotSet]]:
     """
     Replaces placeholders in a template with values from a supplied dictionary.
 
@@ -120,7 +109,7 @@ def apply_values(
     Returns:
         The template with the values applied
     """
-    if isinstance(template, (int, float, bool, type(NotSet), type(None))):
+    if template in (NotSet, None) or isinstance(template, (int, float)):
         return template
     if isinstance(template, str):
         placeholders = find_placeholders(template)
@@ -155,7 +144,7 @@ def apply_values(
 
             return template
     elif isinstance(template, dict):
-        updated_template = {}
+        updated_template: dict[str, Any] = {}
         for key, value in template.items():
             updated_value = apply_values(value, values, remove_notset=remove_notset)
             if updated_value is not NotSet:
@@ -163,22 +152,22 @@ def apply_values(
             elif not remove_notset:
                 updated_template[key] = value
 
-        return updated_template
+        return cast(T, updated_template)
     elif isinstance(template, list):
-        updated_list = []
+        updated_list: list[Any] = []
         for value in template:
             updated_value = apply_values(value, values, remove_notset=remove_notset)
             if updated_value is not NotSet:
                 updated_list.append(updated_value)
-        return updated_list
+        return cast(T, updated_list)
     else:
         raise ValueError(f"Unexpected template type {type(template).__name__!r}")
 
 
 @inject_client
 async def resolve_block_document_references(
-    template: T, client: "PrefectClient" = None
-) -> Union[T, Dict[str, Any]]:
+    template: T, client: Optional["PrefectClient"] = None
+) -> Union[T, dict[str, Any]]:
     """
     Resolve block document references in a template by replacing each reference with
     the data of the block document.
@@ -242,12 +231,17 @@ async def resolve_block_document_references(
     Returns:
         The template with block documents resolved
     """
+    if TYPE_CHECKING:
+        # The @inject_client decorator takes care of providing the client, but
+        # the function signature must mark it as optional to callers.
+        assert client is not None
+
     if isinstance(template, dict):
         block_document_id = template.get("$ref", {}).get("block_document_id")
         if block_document_id:
             block_document = await client.read_block_document(block_document_id)
             return block_document.data
-        updated_template = {}
+        updated_template: dict[str, Any] = {}
         for key, value in template.items():
             updated_value = await resolve_block_document_references(
                 value, client=client
@@ -265,7 +259,7 @@ async def resolve_block_document_references(
             placeholder.type is PlaceholderType.BLOCK_DOCUMENT
             for placeholder in placeholders
         )
-        if len(placeholders) == 0 or not has_block_document_placeholder:
+        if not (placeholders and has_block_document_placeholder):
             return template
         elif (
             len(placeholders) == 1
@@ -274,31 +268,32 @@ async def resolve_block_document_references(
         ):
             # value_keypath will be a list containing a dot path if additional
             # attributes are accessed and an empty list otherwise.
-            block_type_slug, block_document_name, *value_keypath = (
-                list(placeholders)[0]
-                .name.replace(BLOCK_DOCUMENT_PLACEHOLDER_PREFIX, "")
-                .split(".", 2)
-            )
+            [placeholder] = placeholders
+            parts = placeholder.name.replace(
+                BLOCK_DOCUMENT_PLACEHOLDER_PREFIX, ""
+            ).split(".", 2)
+            block_type_slug, block_document_name, *value_keypath = parts
             block_document = await client.read_block_document_by_name(
                 name=block_document_name, block_type_slug=block_type_slug
             )
-            value = block_document.data
+            data = block_document.data
+            value: Union[T, dict[str, Any]] = data
 
             # resolving system blocks to their data for backwards compatibility
-            if len(value) == 1 and "value" in value:
+            if len(data) == 1 and "value" in data:
                 # only resolve the value if the keypath is not already pointing to "value"
-                if len(value_keypath) == 0 or value_keypath[0][:5] != "value":
-                    value = value["value"]
+                if not (value_keypath and value_keypath[0].startswith("value")):
+                    data = value = value["value"]
 
             # resolving keypath/block attributes
-            if len(value_keypath) > 0:
-                value_keypath: str = value_keypath[0]
-                value = get_from_dict(value, value_keypath, default=NotSet)
-                if value is NotSet:
+            if value_keypath:
+                from_dict: Any = get_from_dict(data, value_keypath[0], default=NotSet)
+                if from_dict is NotSet:
                     raise ValueError(
                         f"Invalid template: {template!r}. Could not resolve the"
                         " keypath in the block document data."
                     )
+                value = from_dict
 
             return value
         else:
@@ -311,7 +306,7 @@ async def resolve_block_document_references(
 
 
 @inject_client
-async def resolve_variables(template: T, client: Optional["PrefectClient"] = None):
+async def resolve_variables(template: T, client: Optional["PrefectClient"] = None) -> T:
     """
     Resolve variables in a template by replacing each variable placeholder with the
     value of the variable.
@@ -326,6 +321,11 @@ async def resolve_variables(template: T, client: Optional["PrefectClient"] = Non
     Returns:
         The template with variables resolved
     """
+    if TYPE_CHECKING:
+        # The @inject_client decorator takes care of providing the client, but
+        # the function signature must mark it as optional to callers.
+        assert client is not None
+
     if isinstance(template, str):
         placeholders = find_placeholders(template)
         has_variable_placeholder = any(
@@ -346,7 +346,7 @@ async def resolve_variables(template: T, client: Optional["PrefectClient"] = Non
             if variable is None:
                 return ""
             else:
-                return variable.value
+                return cast(T, variable.value)
         else:
             for full_match, name, placeholder_type in placeholders:
                 if placeholder_type is PlaceholderType.VARIABLE:
@@ -355,7 +355,7 @@ async def resolve_variables(template: T, client: Optional["PrefectClient"] = Non
                     if variable is None:
                         template = template.replace(full_match, "")
                     else:
-                        template = template.replace(full_match, variable.value)
+                        template = template.replace(full_match, str(variable.value))
             return template
     elif isinstance(template, dict):
         return {

prefect/utilities/text.py

@@ -1,5 +1,6 @@
 import difflib
-from typing import Iterable, Optional
+from collections.abc import Iterable
+from typing import Optional
 
 
 def truncated_to(length: int, value: Optional[str]) -> str:

prefect/utilities/timeout.py

@@ -1,6 +1,6 @@
 from asyncio import CancelledError
 from contextlib import contextmanager
-from typing import Optional, Type
+from typing import Optional
 
 from prefect._internal.concurrency.cancellation import (
     cancel_async_after,
@@ -8,7 +8,7 @@ from prefect._internal.concurrency.cancellation import (
 )
 
 
-def fail_if_not_timeout_error(timeout_exc_type: Type[Exception]) -> None:
+def fail_if_not_timeout_error(timeout_exc_type: type[Exception]) -> None:
     if not issubclass(timeout_exc_type, TimeoutError):
         raise ValueError(
             "The `timeout_exc_type` argument must be a subclass of `TimeoutError`."
@@ -17,7 +17,7 @@ def fail_if_not_timeout_error(timeout_exc_type: Type[Exception]) -> None:
 
 @contextmanager
 def timeout_async(
-    seconds: Optional[float] = None, timeout_exc_type: Type[TimeoutError] = TimeoutError
+    seconds: Optional[float] = None, timeout_exc_type: type[TimeoutError] = TimeoutError
 ):
     fail_if_not_timeout_error(timeout_exc_type)
 
@@ -34,7 +34,7 @@ def timeout_async(
 
 @contextmanager
 def timeout(
-    seconds: Optional[float] = None, timeout_exc_type: Type[TimeoutError] = TimeoutError
+    seconds: Optional[float] = None, timeout_exc_type: type[TimeoutError] = TimeoutError
 ):
     fail_if_not_timeout_error(timeout_exc_type)
 

prefect/utilities/urls.py

@@ -2,6 +2,7 @@ import inspect
 import ipaddress
 import socket
 import urllib.parse
+from logging import Logger
 from string import Formatter
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 from urllib.parse import urlparse
@@ -19,7 +20,7 @@ if TYPE_CHECKING:
     from prefect.futures import PrefectFuture
     from prefect.variables import Variable
 
-logger = get_logger("utilities.urls")
+logger: Logger = get_logger("utilities.urls")
 
 # The following objects are excluded from UI URL generation because we lack a
 # directly-addressable URL:
@@ -64,7 +65,7 @@ URLType = Literal["ui", "api"]
 RUN_TYPES = {"flow-run", "task-run"}
 
 
-def validate_restricted_url(url: str):
+def validate_restricted_url(url: str) -> None:
     """
     Validate that the provided URL is safe for outbound requests.  This prevents
     attacks like SSRF (Server Side Request Forgery), where an attacker can make
@@ -123,7 +124,7 @@ def convert_class_to_name(obj: Any) -> str:
 
 def url_for(
     obj: Union[
-        "PrefectFuture",
+        "PrefectFuture[Any]",
         "Block",
         "Variable",
         "Automation",
@@ -163,6 +164,7 @@ def url_for(
         url_for("flow-run", obj_id="123e4567-e89b-12d3-a456-426614174000")
     """
     from prefect.blocks.core import Block
+    from prefect.client.schemas.objects import WorkPool
     from prefect.events.schemas.automations import Automation
     from prefect.events.schemas.events import ReceivedEvent, Resource
     from prefect.futures import PrefectFuture
@@ -228,8 +230,10 @@ def url_for(
         elif name == "block":
             # Blocks are client-side objects whose API representation is a
             # BlockDocument.
-            obj_id = obj._block_document_id
+            obj_id = getattr(obj, "_block_document_id")
         elif name in ("variable", "work-pool"):
+            if TYPE_CHECKING:
+                assert isinstance(obj, (Variable, WorkPool))
             obj_id = obj.name
         elif isinstance(obj, Resource):
             obj_id = obj.id.rpartition(".")[2]
@@ -244,6 +248,7 @@ def url_for(
     url_format = (
         UI_URL_FORMATS.get(name) if url_type == "ui" else API_URL_FORMATS.get(name)
     )
+    assert url_format is not None
 
     if isinstance(obj, ReceivedEvent):
         url = url_format.format(

prefect/utilities/visualization.py

@@ -2,10 +2,12 @@
 Utilities for working with Flow.visualize()
 """
 
+from collections.abc import Coroutine
 from functools import partial
-from typing import Any, List, Optional
+from typing import Any, Literal, Optional, Union, overload
 
-import graphviz
+import graphviz  # type: ignore  # no typing stubs available
+from typing_extensions import Self
 
 from prefect._internal.concurrency.api import from_async
 
@@ -19,7 +21,7 @@ class VisualizationUnsupportedError(Exception):
 
 
 class TaskVizTrackerState:
-    current = None
+    current: Optional["TaskVizTracker"] = None
 
 
 class GraphvizImportError(Exception):
@@ -30,16 +32,36 @@ class GraphvizExecutableNotFoundError(Exception):
     pass
 
 
-def get_task_viz_tracker():
+def get_task_viz_tracker() -> Optional["TaskVizTracker"]:
     return TaskVizTrackerState.current
 
 
+@overload
+def track_viz_task(
+    is_async: Literal[True],
+    task_name: str,
+    parameters: dict[str, Any],
+    viz_return_value: Optional[Any] = None,
+) -> Coroutine[Any, Any, Any]:
+    ...
+
+
+@overload
+def track_viz_task(
+    is_async: Literal[False],
+    task_name: str,
+    parameters: dict[str, Any],
+    viz_return_value: Optional[Any] = None,
+) -> Any:
+    ...
+
+
 def track_viz_task(
     is_async: bool,
     task_name: str,
-    parameters: dict,
+    parameters: dict[str, Any],
     viz_return_value: Optional[Any] = None,
-):
+) -> Union[Coroutine[Any, Any, Any], Any]:
     """Return a result if sync otherwise return a coroutine that returns the result"""
     if is_async:
         return from_async.wait_for_call_in_loop_thread(
@@ -50,14 +72,14 @@ def track_viz_task(
 
 
 def _track_viz_task(
-    task_name,
-    parameters,
-    viz_return_value=None,
+    task_name: str,
+    parameters: dict[str, Any],
+    viz_return_value: Optional[Any] = None,
 ) -> Any:
     task_run_tracker = get_task_viz_tracker()
     if task_run_tracker:
-        upstream_tasks = []
-        for k, v in parameters.items():
+        upstream_tasks: list[VizTask] = []
+        for _, v in parameters.items():
             if isinstance(v, VizTask):
                 upstream_tasks.append(v)
             # if it's an object that we've already seen,
@@ -85,19 +107,19 @@ class VizTask:
     def __init__(
         self,
         name: str,
-        upstream_tasks: Optional[List["VizTask"]] = None,
+        upstream_tasks: Optional[list["VizTask"]] = None,
     ):
         self.name = name
-        self.upstream_tasks = upstream_tasks if upstream_tasks else []
+        self.upstream_tasks: list[VizTask] = upstream_tasks if upstream_tasks else []
 
 
 class TaskVizTracker:
     def __init__(self):
-        self.tasks = []
-        self.dynamic_task_counter = {}
-        self.object_id_to_task = {}
+        self.tasks: list[VizTask] = []
+        self.dynamic_task_counter: dict[str, int] = {}
+        self.object_id_to_task: dict[int, VizTask] = {}
 
-    def add_task(self, task: VizTask):
+    def add_task(self, task: VizTask) -> None:
         if task.name not in self.dynamic_task_counter:
             self.dynamic_task_counter[task.name] = 0
         else:
@@ -106,11 +128,11 @@ class TaskVizTracker:
         task.name = f"{task.name}-{self.dynamic_task_counter[task.name]}"
         self.tasks.append(task)
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         TaskVizTrackerState.current = self
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         TaskVizTrackerState.current = None
 
     def link_viz_return_value_to_viz_task(
@@ -129,7 +151,7 @@ class TaskVizTracker:
         self.object_id_to_task[id(viz_return_value)] = viz_task
 
 
-def build_task_dependencies(task_run_tracker: TaskVizTracker):
+def build_task_dependencies(task_run_tracker: TaskVizTracker) -> graphviz.Digraph:
     """
     Constructs a Graphviz directed graph object that represents the dependencies
     between tasks in the given TaskVizTracker.
@@ -151,9 +173,9 @@ def build_task_dependencies(task_run_tracker: TaskVizTracker):
     try:
         g = graphviz.Digraph()
         for task in task_run_tracker.tasks:
-            g.node(task.name)
+            g.node(task.name)  # type: ignore[reportUnknownMemberType]
             for upstream in task.upstream_tasks:
-                g.edge(upstream.name, task.name)
+                g.edge(upstream.name, task.name)  # type: ignore[reportUnknownMemberType]
         return g
     except ImportError as exc:
         raise GraphvizImportError from exc
@@ -166,7 +188,7 @@ def build_task_dependencies(task_run_tracker: TaskVizTracker):
         )
 
 
-def visualize_task_dependencies(graph: graphviz.Digraph, flow_run_name: str):
+def visualize_task_dependencies(graph: graphviz.Digraph, flow_run_name: str) -> None:
     """
     Renders and displays a Graphviz directed graph representing task dependencies.
 
@@ -184,7 +206,7 @@ def visualize_task_dependencies(graph: graphviz.Digraph, flow_run_name: str):
       specifying a `viz_return_value`.
     """
     try:
-        graph.render(filename=flow_run_name, view=True, format="png", cleanup=True)
+        graph.render(filename=flow_run_name, view=True, format="png", cleanup=True)  # type: ignore[reportUnknownMemberType]
     except graphviz.backend.ExecutableNotFound as exc:
         msg = (
             "It appears you do not have Graphviz installed, or it is not on your "