prefect-client 3.1.4__py3-none-any.whl → 3.1.6__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 "

prefect/variables.py

@@ -73,17 +73,18 @@ class Variable(BaseModel):
                 raise ValueError(
                     f"Variable {name!r} already exists. Use `overwrite=True` to update it."
                 )
-            await client.update_variable(variable=VariableUpdate(**var_dict))
+            await client.update_variable(
+                variable=VariableUpdate.model_validate(var_dict)
+            )
             variable = await client.read_variable_by_name(name)
-            var_dict = {
-                "name": variable.name,
-                "value": variable.value,
-                "tags": variable.tags or [],
-            }
+            for key in var_dict.keys():
+                var_dict.update({key: getattr(variable, key)})
         else:
-            await client.create_variable(variable=VariableCreate(**var_dict))
+            await client.create_variable(
+                variable=VariableCreate.model_validate(var_dict)
+            )
 
-        return cls(**var_dict)
+        return cls.model_validate(var_dict)
 
     @classmethod
     @sync_compatible

prefect/workers/base.py

@@ -53,6 +53,7 @@ from prefect.states import (
     Pending,
     exception_to_failed_state,
 )
+from prefect.types import KeyValueLabels
 from prefect.utilities.dispatch import get_registry_for_type, register_base_type
 from prefect.utilities.engine import propose_state
 from prefect.utilities.services import critical_service_loop
@@ -637,8 +638,9 @@ class BaseWorker(abc.ABC):
         await self._exit_stack.enter_async_context(self._runs_task_group)
 
         if self._client.server_type == ServerType.CLOUD:
-            self._cloud_client = get_cloud_client()
-            await self._exit_stack.enter_async_context(self._cloud_client)
+            self._cloud_client = await self._exit_stack.enter_async_context(
+                get_cloud_client()
+            )
 
         self.is_setup = True
 
@@ -987,7 +989,6 @@ class BaseWorker(abc.ABC):
         try:
             configuration = await self._get_configuration(flow_run)
             submitted_event = self._emit_flow_run_submitted_event(configuration)
-            await self._give_worker_labels_to_flow_run(flow_run.id)
             result = await self.run(
                 flow_run=flow_run,
                 task_status=task_status,
@@ -1221,13 +1222,20 @@ class BaseWorker(abc.ABC):
         Give this worker's identifying labels to the specified flow run.
         """
         if self._cloud_client:
-            await self._cloud_client.update_flow_run_labels(
-                flow_run_id,
-                {
-                    "prefect.worker.name": self.name,
-                    "prefect.worker.type": self.type,
-                },
-            )
+            labels: KeyValueLabels = {
+                "prefect.worker.name": self.name,
+                "prefect.worker.type": self.type,
+            }
+
+            if self._work_pool:
+                labels.update(
+                    {
+                        "prefect.work-pool.name": self._work_pool.name,
+                        "prefect.work-pool.id": str(self._work_pool.id),
+                    }
+                )
+
+            await self._cloud_client.update_flow_run_labels(flow_run_id, labels)
 
     async def __aenter__(self):
         self._logger.debug("Entering worker context...")

{prefect_client-3.1.4.dist-info → prefect_client-3.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.1.4
+Version: 3.1.6
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.
@@ -45,7 +45,7 @@ Requires-Dist: packaging<24.3,>=21.3
 Requires-Dist: pathspec>=0.8.0
 Requires-Dist: pendulum<4,>=3.0.0
 Requires-Dist: prometheus-client>=0.20.0
-Requires-Dist: pydantic<3.0.0,>=2.7
+Requires-Dist: pydantic!=2.10.0,<3.0.0,>=2.7
 Requires-Dist: pydantic-core<3.0.0,>=2.12.0
 Requires-Dist: pydantic-extra-types<3.0.0,>=2.8.2
 Requires-Dist: pydantic-settings>2.2.1
@@ -79,8 +79,6 @@ Requires-Dist: apprise<2.0.0,>=1.1.0; extra == "notifications"
     <br>
     <a href="https://prefect.io/slack" alt="Slack">
         <img src="https://img.shields.io/badge/slack-join_community-red.svg?color=0052FF&labelColor=090422&logo=slack" /></a>
-    <a href="https://discourse.prefect.io/" alt="Discourse">
-        <img src="https://img.shields.io/badge/discourse-browse_forum-red.svg?color=0052FF&labelColor=090422&logo=discourse" /></a>
     <a href="https://www.youtube.com/c/PrefectIO/" alt="YouTube">
         <img src="https://img.shields.io/badge/youtube-watch_videos-red.svg?color=0052FF&labelColor=090422&logo=youtube" /></a>
 </p>
@@ -167,7 +165,9 @@ Start with our [friendly tutorial](https://docs.prefect.io/tutorials) or explore
 
 ## Join the community
 
-Prefect is made possible by the fastest growing community of thousands of friendly data engineers. Join us in building a new kind of workflow system. The [Prefect Slack community](https://prefect.io/slack) is a fantastic place to learn more about Prefect, ask questions, or get help with workflow design. All community forums, including code contributions, issue discussions, and slack messages are subject to our [Code of Conduct](https://discourse.prefect.io/faq).
+Prefect is made possible by the fastest growing community of thousands of friendly data engineers. Join us in building a new kind of workflow system. 
+The [Prefect Slack community](https://prefect.io/slack) is a fantastic place to learn more about Prefect, ask questions, or get help with workflow design. 
+All community forums, including code contributions, issue discussions, and Slack messages are subject to our [Code of Conduct](https://github.com/PrefectHQ/prefect/blob/main/CODE_OF_CONDUCT.md).
 
 ## Contribute