prefect-client 3.1.9__py3-none-any.whl → 3.1.10__py3-none-any.whl

prefect/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-20T16:33:15-0500",
+ "date": "2024-12-24T22:58:40-0800",
  "dirty": true,
  "error": null,
- "full-revisionid": "e1fe79439b319ce5478aa5f8bd4e836ed4f8a0f5",
- "version": "3.1.9"
+ "full-revisionid": "b11b56fc6143df2bd4c7018f7a974dc656fe9aa3",
+ "version": "3.1.10"
 }
 '''  # END VERSION_JSON
 

prefect/artifacts.py

@@ -2,19 +2,20 @@
 Interface for creating and reading artifacts.
 """
 
-from __future__ import annotations
-
+import asyncio
 import json  # noqa: I001
 import math
 import warnings
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 from uuid import UUID
 
+from typing_extensions import Self
+
 from prefect.client.schemas.actions import ArtifactCreate as ArtifactRequest
 from prefect.client.schemas.actions import ArtifactUpdate
 from prefect.client.schemas.filters import ArtifactFilter, ArtifactFilterKey
 from prefect.client.schemas.sorting import ArtifactSort
-from prefect.client.utilities import get_or_create_client, inject_client
+from prefect.client.utilities import get_or_create_client
 from prefect.logging.loggers import get_logger
 from prefect.utilities.asyncutils import sync_compatible
 from prefect.utilities.context import get_task_and_flow_run_ids
@@ -22,8 +23,6 @@ from prefect.utilities.context import get_task_and_flow_run_ids
 logger = get_logger("artifacts")
 
 if TYPE_CHECKING:
-    from typing_extensions import Self
-
     from prefect.client.orchestration import PrefectClient
     from prefect.client.schemas.objects import Artifact as ArtifactResponse
 
@@ -43,7 +42,7 @@ class Artifact(ArtifactRequest):
 
     @sync_compatible
     async def create(
-        self: "Self",
+        self: Self,
         client: Optional["PrefectClient"] = None,
     ) -> "ArtifactResponse":
         """
@@ -95,16 +94,15 @@ class Artifact(ArtifactRequest):
             (ArtifactResponse, optional): The artifact (if found).
         """
         client, _ = get_or_create_client(client)
-        return next(
-            iter(
-                await client.read_artifacts(
-                    limit=1,
-                    sort=ArtifactSort.UPDATED_DESC,
-                    artifact_filter=ArtifactFilter(key=ArtifactFilterKey(any_=[key])),
-                )
+        filter_key_value = None if key is None else [key]
+        artifacts = await client.read_artifacts(
+            limit=1,
+            sort=ArtifactSort.UPDATED_DESC,
+            artifact_filter=ArtifactFilter(
+                key=ArtifactFilterKey(any_=filter_key_value)
             ),
-            None,
         )
+        return None if not artifacts else artifacts[0]
 
     @classmethod
     @sync_compatible
@@ -112,10 +110,10 @@ class Artifact(ArtifactRequest):
         cls,
         key: Optional[str] = None,
         description: Optional[str] = None,
-        data: Optional[Union[Dict[str, Any], Any]] = None,
+        data: Optional[Union[dict[str, Any], Any]] = None,
         client: Optional["PrefectClient"] = None,
         **kwargs: Any,
-    ) -> Tuple["ArtifactResponse", bool]:
+    ) -> tuple["ArtifactResponse", bool]:
         """
         A method to get or create an artifact.
 
@@ -128,18 +126,20 @@ class Artifact(ArtifactRequest):
         Returns:
             (ArtifactResponse): The artifact, either retrieved or created.
         """
-        artifact = await cls.get(key, client)
+        artifact_coro = cls.get(key, client)
+        if TYPE_CHECKING:
+            assert asyncio.iscoroutine(artifact_coro)
+        artifact = await artifact_coro
         if artifact:
             return artifact, False
-        else:
-            return (
-                await cls(key=key, description=description, data=data, **kwargs).create(
-                    client
-                ),
-                True,
-            )
 
-    async def format(self) -> Optional[Union[Dict[str, Any], Any]]:
+        new_artifact = cls(key=key, description=description, data=data, **kwargs)
+        create_coro = new_artifact.create(client)
+        if TYPE_CHECKING:
+            assert asyncio.iscoroutine(create_coro)
+        return await create_coro, True
+
+    async def format(self) -> Optional[Union[dict[str, Any], Any]]:
         return json.dumps(self.data)
 
 
@@ -165,13 +165,13 @@ class MarkdownArtifact(Artifact):
 
 
 class TableArtifact(Artifact):
-    table: Union[Dict[str, List[Any]], List[Dict[str, Any]], List[List[Any]]]
+    table: Union[dict[str, list[Any]], list[dict[str, Any]], list[list[Any]]]
     type: Optional[str] = "table"
 
     @classmethod
     def _sanitize(
-        cls, item: Union[Dict[str, Any], List[Any], float]
-    ) -> Union[Dict[str, Any], List[Any], int, float, None]:
+        cls, item: Union[dict[str, Any], list[Any], float]
+    ) -> Union[dict[str, Any], list[Any], int, float, None]:
         """
         Sanitize NaN values in a given item.
         The item can be a dict, list or float.
@@ -230,39 +230,6 @@ class ImageArtifact(Artifact):
         return self.image_url
 
 
-@inject_client
-async def _create_artifact(
-    type: str,
-    key: Optional[str] = None,
-    description: Optional[str] = None,
-    data: Optional[Union[Dict[str, Any], Any]] = None,
-    client: Optional["PrefectClient"] = None,
-) -> UUID:
-    """
-    Helper function to create an artifact.
-
-    Arguments:
-        type: A string identifying the type of artifact.
-        key: A user-provided string identifier.
-          The key must only contain lowercase letters, numbers, and dashes.
-        description: A user-specified description of the artifact.
-        data: A JSON payload that allows for a result to be retrieved.
-        client: The PrefectClient
-
-    Returns:
-        - The table artifact ID.
-    """
-
-    artifact = await Artifact(
-        type=type,
-        key=key,
-        description=description,
-        data=data,
-    ).create(client)
-
-    return artifact.id
-
-
 @sync_compatible
 async def create_link_artifact(
     link: str,
@@ -286,12 +253,16 @@ async def create_link_artifact(
     Returns:
         The table artifact ID.
     """
-    artifact = await LinkArtifact(
+    new_artifact = LinkArtifact(
         key=key,
         description=description,
         link=link,
         link_text=link_text,
-    ).create(client)
+    )
+    create_coro = new_artifact.create(client)
+    if TYPE_CHECKING:
+        assert asyncio.iscoroutine(create_coro)
+    artifact = await create_coro
 
     return artifact.id
 
@@ -315,18 +286,22 @@ async def create_markdown_artifact(
     Returns:
         The table artifact ID.
     """
-    artifact = await MarkdownArtifact(
+    new_artifact = MarkdownArtifact(
         key=key,
         description=description,
         markdown=markdown,
-    ).create()
+    )
+    create_coro = new_artifact.create()
+    if TYPE_CHECKING:
+        assert asyncio.iscoroutine(create_coro)
+    artifact = await create_coro
 
     return artifact.id
 
 
 @sync_compatible
 async def create_table_artifact(
-    table: Union[Dict[str, List[Any]], List[Dict[str, Any]], List[List[Any]]],
+    table: Union[dict[str, list[Any]], list[dict[str, Any]], list[list[Any]]],
     key: Optional[str] = None,
     description: Optional[str] = None,
 ) -> UUID:
@@ -344,11 +319,15 @@ async def create_table_artifact(
         The table artifact ID.
     """
 
-    artifact = await TableArtifact(
+    new_artifact = TableArtifact(
         key=key,
         description=description,
         table=table,
-    ).create()
+    )
+    create_coro = new_artifact.create()
+    if TYPE_CHECKING:
+        assert asyncio.iscoroutine(create_coro)
+    artifact = await create_coro
 
     return artifact.id
 
@@ -373,11 +352,15 @@ async def create_progress_artifact(
         The progress artifact ID.
     """
 
-    artifact = await ProgressArtifact(
+    new_artifact = ProgressArtifact(
         key=key,
         description=description,
         progress=progress,
-    ).create()
+    )
+    create_coro = new_artifact.create()
+    if TYPE_CHECKING:
+        assert asyncio.iscoroutine(create_coro)
+    artifact = await create_coro
 
     return artifact.id
 
@@ -387,7 +370,7 @@ async def update_progress_artifact(
     artifact_id: UUID,
     progress: float,
     description: Optional[str] = None,
-    client: Optional[PrefectClient] = None,
+    client: Optional["PrefectClient"] = None,
 ) -> UUID:
     """
     Update a progress artifact.
@@ -444,10 +427,14 @@ async def create_image_artifact(
         The image artifact ID.
     """
 
-    artifact = await ImageArtifact(
+    new_artifact = ImageArtifact(
         key=key,
         description=description,
         image_url=image_url,
-    ).create()
+    )
+    create_coro = new_artifact.create()
+    if TYPE_CHECKING:
+        assert asyncio.iscoroutine(create_coro)
+    artifact = await create_coro
 
     return artifact.id

prefect/automations.py

@@ -1,4 +1,4 @@
-from typing import Optional, Type
+from typing import TYPE_CHECKING, Optional, overload
 from uuid import UUID
 
 from pydantic import Field
@@ -112,17 +112,28 @@ class Automation(AutomationCore):
         auto.name = "new name"
         auto.update()
         """
+        assert self.id is not None
         async with get_client() as client:
             automation = AutomationCore(
                 **self.model_dump(exclude={"id", "owner_resource"})
             )
             await client.update_automation(automation_id=self.id, automation=automation)
 
+    @overload
+    @classmethod
+    async def read(cls, id: UUID, name: Optional[str] = ...) -> Self:
+        ...
+
+    @overload
+    @classmethod
+    async def read(cls, id: None = None, name: str = ...) -> Optional[Self]:
+        ...
+
     @classmethod
     @sync_compatible
     async def read(
-        cls: Type[Self], id: Optional[UUID] = None, name: Optional[str] = None
-    ) -> Self:
+        cls, id: Optional[UUID] = None, name: Optional[str] = None
+    ) -> Optional[Self]:
         """
         Read an automation by ID or name.
         automation = Automation.read(name="woodchonk")
@@ -145,13 +156,13 @@ class Automation(AutomationCore):
                     raise
                 if automation is None:
                     raise ValueError(f"Automation with ID {id!r} not found")
-                return Automation(**automation.model_dump())
+                return cls(**automation.model_dump())
             else:
+                if TYPE_CHECKING:
+                    assert name is not None
                 automation = await client.read_automations_by_name(name=name)
                 if len(automation) > 0:
-                    return (
-                        Automation(**automation[0].model_dump()) if automation else None
-                    )
+                    return cls(**automation[0].model_dump()) if automation else None
                 else:
                     raise ValueError(f"Automation with name {name!r} not found")
 
@@ -161,6 +172,9 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.delete()
         """
+        if self.id is None:
+            raise ValueError("Can't delete an automation without an id")
+
         async with get_client() as client:
             try:
                 await client.delete_automation(self.id)
@@ -177,6 +191,9 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.disable()
         """
+        if self.id is None:
+            raise ValueError("Can't disable an automation without an id")
+
         async with get_client() as client:
             try:
                 await client.pause_automation(self.id)
@@ -193,6 +210,9 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.enable()
         """
+        if self.id is None:
+            raise ValueError("Can't enable an automation without an id")
+
         async with get_client() as client:
             try:
                 await client.resume_automation(self.id)

prefect/context.py

@@ -9,21 +9,10 @@ For more user-accessible information about the current run, see [`prefect.runtim
 import os
 import sys
 import warnings
+from collections.abc import AsyncGenerator, Generator, Mapping
 from contextlib import ExitStack, asynccontextmanager, contextmanager
 from contextvars import ContextVar, Token
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    AsyncGenerator,
-    Dict,
-    Generator,
-    Mapping,
-    Optional,
-    Set,
-    Type,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
 from typing_extensions import Self
@@ -64,7 +53,7 @@ if TYPE_CHECKING:
 GLOBAL_SETTINGS_CONTEXT = None  # type: ignore
 
 
-def serialize_context() -> Dict[str, Any]:
+def serialize_context() -> dict[str, Any]:
     """
     Serialize the current context for use in a remote execution environment.
     """
@@ -84,7 +73,7 @@ def serialize_context() -> Dict[str, Any]:
 
 @contextmanager
 def hydrated_context(
-    serialized_context: Optional[Dict[str, Any]] = None,
+    serialized_context: Optional[dict[str, Any]] = None,
     client: Union[PrefectClient, SyncPrefectClient, None] = None,
 ):
     with ExitStack() as stack:
@@ -148,7 +137,7 @@ class ContextModel(BaseModel):
         self._token = None
 
     @classmethod
-    def get(cls: Type[Self]) -> Optional[Self]:
+    def get(cls: type[Self]) -> Optional[Self]:
         """Get the current context instance"""
         return cls.__var__.get(None)
 
@@ -173,7 +162,7 @@ class ContextModel(BaseModel):
         new._token = None
         return new
 
-    def serialize(self, include_secrets: bool = True) -> Dict[str, Any]:
+    def serialize(self, include_secrets: bool = True) -> dict[str, Any]:
         """
         Serialize the context model to a dictionary that can be pickled with cloudpickle.
         """
@@ -314,10 +303,10 @@ class RunContext(ContextModel):
         start_client_metrics_server()
 
     start_time: DateTime = Field(default_factory=lambda: DateTime.now("UTC"))
-    input_keyset: Optional[Dict[str, Dict[str, str]]] = None
+    input_keyset: Optional[dict[str, dict[str, str]]] = None
     client: Union[PrefectClient, SyncPrefectClient]
 
-    def serialize(self: Self, include_secrets: bool = True) -> Dict[str, Any]:
+    def serialize(self: Self, include_secrets: bool = True) -> dict[str, Any]:
         return self.model_dump(
             include={"start_time", "input_keyset"},
             exclude_unset=True,
@@ -344,7 +333,7 @@ class EngineContext(RunContext):
     flow_run: Optional[FlowRun] = None
     task_runner: TaskRunner[Any]
     log_prints: bool = False
-    parameters: Optional[Dict[str, Any]] = None
+    parameters: Optional[dict[str, Any]] = None
 
     # Flag signaling if the flow run context has been serialized and sent
     # to remote infrastructure.
@@ -355,10 +344,10 @@ class EngineContext(RunContext):
     persist_result: bool = Field(default_factory=get_default_persist_setting)
 
     # Counter for task calls allowing unique
-    task_run_dynamic_keys: Dict[str, Union[str, int]] = Field(default_factory=dict)
+    task_run_dynamic_keys: dict[str, Union[str, int]] = Field(default_factory=dict)
 
     # Counter for flow pauses
-    observed_flow_pauses: Dict[str, int] = Field(default_factory=dict)
+    observed_flow_pauses: dict[str, int] = Field(default_factory=dict)
 
     # Tracking for result from task runs in this flow run for dependency tracking
     # Holds the ID of the object returned by the task run and task run state
@@ -369,7 +358,7 @@ class EngineContext(RunContext):
 
     __var__: ContextVar[Self] = ContextVar("flow_run")
 
-    def serialize(self: Self, include_secrets: bool = True) -> Dict[str, Any]:
+    def serialize(self: Self, include_secrets: bool = True) -> dict[str, Any]:
         return self.model_dump(
             include={
                 "flow_run",
@@ -403,7 +392,7 @@ class TaskRunContext(RunContext):
     task: "Task[Any, Any]"
     task_run: TaskRun
     log_prints: bool = False
-    parameters: Dict[str, Any]
+    parameters: dict[str, Any]
 
     # Result handling
     result_store: ResultStore
@@ -411,7 +400,7 @@ class TaskRunContext(RunContext):
 
     __var__ = ContextVar("task_run")
 
-    def serialize(self: Self, include_secrets: bool = True) -> Dict[str, Any]:
+    def serialize(self: Self, include_secrets: bool = True) -> dict[str, Any]:
         return self.model_dump(
             include={
                 "task_run",
@@ -437,7 +426,7 @@ class TagsContext(ContextModel):
         current_tags: A set of current tags in the context
     """
 
-    current_tags: Set[str] = Field(default_factory=set)
+    current_tags: set[str] = Field(default_factory=set)
 
     @classmethod
     def get(cls) -> "TagsContext":
@@ -512,7 +501,7 @@ def get_settings_context() -> SettingsContext:
 
 
 @contextmanager
-def tags(*new_tags: str) -> Generator[Set[str], None, None]:
+def tags(*new_tags: str) -> Generator[set[str], None, None]:
     """
     Context manager to add tags to flow and task run calls.
 

prefect/flows.py

@@ -37,7 +37,6 @@ from typing import (
 from uuid import UUID
 
 import pydantic
-from fastapi.encoders import jsonable_encoder
 from pydantic.v1 import BaseModel as V1BaseModel
 from pydantic.v1.decorator import ValidatedFunction as V1ValidatedFunction
 from pydantic.v1.errors import ConfigError  # TODO
@@ -613,6 +612,8 @@ class Flow(Generic[P, R]):
                 serialized_parameters[key] = f"<{type(value).__name__}>"
                 continue
             try:
+                from fastapi.encoders import jsonable_encoder
+
                 serialized_parameters[key] = jsonable_encoder(value)
             except (TypeError, ValueError):
                 logger.debug(

prefect/futures.py

@@ -1,17 +1,15 @@
 import abc
 import asyncio
-import collections
 import concurrent.futures
 import threading
 import uuid
 from collections.abc import Generator, Iterator
 from functools import partial
-from typing import Any, Callable, Generic, List, Optional, Set, Union, cast
+from typing import Any, Callable, Generic, Optional, Union
 
-from typing_extensions import TypeVar
+from typing_extensions import NamedTuple, Self, TypeVar
 
 from prefect.client.orchestration import get_client
-from prefect.client.schemas.objects import TaskRun
 from prefect.exceptions import ObjectNotFound
 from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.states import Pending, State
@@ -50,7 +48,7 @@ class PrefectFuture(abc.ABC, Generic[R]):
             return self._final_state
         client = get_client(sync_client=True)
         try:
-            task_run = cast(TaskRun, client.read_task_run(task_run_id=self.task_run_id))
+            task_run = client.read_task_run(task_run_id=self.task_run_id)
         except ObjectNotFound:
             # We'll be optimistic and assume this task will eventually start
             # TODO: Consider using task run events to wait for the task to start
@@ -92,7 +90,7 @@ class PrefectFuture(abc.ABC, Generic[R]):
         """
 
     @abc.abstractmethod
-    def add_done_callback(self, fn):
+    def add_done_callback(self, fn: Callable[["PrefectFuture[R]"], None]):
         """
         Add a callback to be run when the future completes or is cancelled.
 
@@ -102,13 +100,17 @@ class PrefectFuture(abc.ABC, Generic[R]):
         ...
 
 
-class PrefectWrappedFuture(PrefectFuture, abc.ABC, Generic[R, F]):
+class PrefectWrappedFuture(PrefectFuture[R], abc.ABC, Generic[R, F]):
     """
     A Prefect future that wraps another future object.
+
+    Type Parameters:
+        R: The return type of the future
+        F: The type of the wrapped future
     """
 
     def __init__(self, task_run_id: uuid.UUID, wrapped_future: F):
-        self._wrapped_future = wrapped_future
+        self._wrapped_future: F = wrapped_future
         super().__init__(task_run_id)
 
     @property
@@ -116,10 +118,11 @@ class PrefectWrappedFuture(PrefectFuture, abc.ABC, Generic[R, F]):
         """The underlying future object wrapped by this Prefect future"""
         return self._wrapped_future
 
-    def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]):
+    def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]) -> None:
+        """Add a callback to be executed when the future completes."""
         if not self._final_state:
 
-            def call_with_self(future):
+            def call_with_self(future: F):
                 """Call the callback with self as the argument, this is necessary to ensure we remove the future from the pending set"""
                 fn(self)
 
@@ -128,7 +131,7 @@ class PrefectWrappedFuture(PrefectFuture, abc.ABC, Generic[R, F]):
         fn(self)
 
 
-class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future]):
+class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future[R]]):
     """
     A Prefect future that wraps a concurrent.futures.Future. This future is used
     when the task run is submitted to a ThreadPoolExecutor.
@@ -193,7 +196,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
     any task run scheduled in Prefect's API.
     """
 
-    done_callbacks: List[Callable[[PrefectFuture[R]], None]] = []
+    done_callbacks: list[Callable[[PrefectFuture[R]], None]] = []
     waiter = None
 
     def wait(self, timeout: Optional[float] = None) -> None:
@@ -270,7 +273,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
                 return
             TaskRunWaiter.add_done_callback(self._task_run_id, partial(fn, self))
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         if not isinstance(other, PrefectDistributedFuture):
             return False
         return self.task_run_id == other.task_run_id
@@ -279,7 +282,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
         return hash(self.task_run_id)
 
 
-class PrefectFutureList(list, Iterator, Generic[F]):
+class PrefectFutureList(list[PrefectFuture[R]], Iterator[PrefectFuture[R]]):
     """
     A list of Prefect futures.
 
@@ -298,10 +301,10 @@ class PrefectFutureList(list, Iterator, Generic[F]):
         wait(self, timeout=timeout)
 
     def result(
-        self: "PrefectFutureList[R]",
+        self: Self,
         timeout: Optional[float] = None,
         raise_on_failure: bool = True,
-    ) -> List[R]:
+    ) -> list[R]:
         """
         Get the results of all task runs associated with the futures in the list.
 
@@ -331,21 +334,22 @@ class PrefectFutureList(list, Iterator, Generic[F]):
 
 
 def as_completed(
-    futures: List[PrefectFuture[R]], timeout: Optional[float] = None
+    futures: list[PrefectFuture[R]], timeout: Optional[float] = None
 ) -> Generator[PrefectFuture[R], None]:
-    unique_futures: Set[PrefectFuture[R]] = set(futures)
+    unique_futures: set[PrefectFuture[R]] = set(futures)
     total_futures = len(unique_futures)
+    pending = unique_futures
     try:
         with timeout_context(timeout):
-            done = {f for f in unique_futures if f._final_state}
+            done = {f for f in unique_futures if f._final_state}  # type: ignore[privateUsage]
             pending = unique_futures - done
             yield from done
 
             finished_event = threading.Event()
             finished_lock = threading.Lock()
-            finished_futures = []
+            finished_futures: list[PrefectFuture[R]] = []
 
-            def add_to_done(future):
+            def add_to_done(future: PrefectFuture[R]):
                 with finished_lock:
                     finished_futures.append(future)
                     finished_event.set()
@@ -370,10 +374,19 @@ def as_completed(
         )
 
 
-DoneAndNotDoneFutures = collections.namedtuple("DoneAndNotDoneFutures", "done not_done")
+class DoneAndNotDoneFutures(NamedTuple, Generic[R]):
+    """A named 2-tuple of sets.
+
+    multiple inheritance supported in 3.11+, use typing_extensions.NamedTuple
+    """
+
+    done: set[PrefectFuture[R]]
+    not_done: set[PrefectFuture[R]]
 
 
-def wait(futures: List[PrefectFuture[R]], timeout=None) -> DoneAndNotDoneFutures:
+def wait(
+    futures: list[PrefectFuture[R]], timeout: Optional[float] = None
+) -> DoneAndNotDoneFutures[R]:
     """
     Wait for the futures in the given sequence to complete.
 
@@ -431,9 +444,11 @@ def resolve_futures_to_states(
 
     Unsupported object types will be returned without modification.
     """
-    futures: Set[PrefectFuture[R]] = set()
+    futures: set[PrefectFuture[R]] = set()
 
-    def _collect_futures(futures, expr, context):
+    def _collect_futures(
+        futures: set[PrefectFuture[R]], expr: Any, context: Any
+    ) -> Union[PrefectFuture[R], Any]:
         # Expressions inside quotes should not be traversed
         if isinstance(context.get("annotation"), quote):
             raise StopVisiting()
@@ -455,14 +470,14 @@ def resolve_futures_to_states(
         return expr
 
     # Get final states for each future
-    states = []
+    states: list[State] = []
     for future in futures:
         future.wait()
         states.append(future.state)
 
     states_by_future = dict(zip(futures, states))
 
-    def replace_futures_with_states(expr, context):
+    def replace_futures_with_states(expr: Any, context: Any) -> Any:
         # Expressions inside quotes should not be modified
         if isinstance(context.get("annotation"), quote):
             raise StopVisiting()

prefect/locking/filesystem.py

@@ -1,6 +1,7 @@
 import time
+from logging import Logger
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Optional
 
 import anyio
 import pendulum
@@ -11,7 +12,7 @@ from prefect.logging.loggers import get_logger
 
 from .protocol import LockManager
 
-logger = get_logger(__name__)
+logger: Logger = get_logger(__name__)
 
 
 class _LockInfo(TypedDict):
@@ -37,11 +38,11 @@ class FileSystemLockManager(LockManager):
         lock_files_directory: the directory where lock files are stored
     """
 
-    def __init__(self, lock_files_directory: Path):
-        self.lock_files_directory = lock_files_directory.expanduser().resolve()
-        self._locks: Dict[str, _LockInfo] = {}
+    def __init__(self, lock_files_directory: Path) -> None:
+        self.lock_files_directory: Path = lock_files_directory.expanduser().resolve()
+        self._locks: dict[str, _LockInfo] = {}
 
-    def _ensure_lock_files_directory_exists(self):
+    def _ensure_lock_files_directory_exists(self) -> None:
         self.lock_files_directory.mkdir(parents=True, exist_ok=True)
 
     def _lock_path_for_key(self, key: str) -> Path:
@@ -49,7 +50,7 @@ class FileSystemLockManager(LockManager):
             return lock_info["path"]
         return self.lock_files_directory.joinpath(key).with_suffix(".lock")
 
-    def _get_lock_info(self, key: str, use_cache=True) -> Optional[_LockInfo]:
+    def _get_lock_info(self, key: str, use_cache: bool = True) -> Optional[_LockInfo]:
         if use_cache:
             if (lock_info := self._locks.get(key)) is not None:
                 return lock_info

prefect/locking/memory.py

@@ -1,6 +1,8 @@
 import asyncio
 import threading
-from typing import Dict, Optional, TypedDict
+from typing import Any, Optional, TypedDict
+
+from typing_extensions import Self
 
 from .protocol import LockManager
 
@@ -30,14 +32,14 @@ class MemoryLockManager(LockManager):
 
     _instance = None
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: Any, **kwargs: Any) -> Self:
         if cls._instance is None:
             cls._instance = super().__new__(cls)
         return cls._instance
 
     def __init__(self):
         self._locks_dict_lock = threading.Lock()
-        self._locks: Dict[str, _LockInfo] = {}
+        self._locks: dict[str, _LockInfo] = {}
 
     def _expire_lock(self, key: str):
         """

prefect/locking/protocol.py

@@ -57,7 +57,7 @@ class LockManager(Protocol):
         """
         ...
 
-    def release_lock(self, key: str, holder: str):
+    def release_lock(self, key: str, holder: str) -> None:
         """
         Releases the lock on the corresponding transaction record.
 

prefect/plugins.py

@@ -9,15 +9,15 @@ Currently supported entrypoints:
 """
 
 from types import ModuleType
-from typing import Any, Dict, Union
+from typing import Any, Union
 
 import prefect.settings
 from prefect.utilities.compat import EntryPoints, entry_points
 
-COLLECTIONS: Union[None, Dict[str, Union[ModuleType, Exception]]] = None
+_collections: Union[None, dict[str, Union[ModuleType, Exception]]] = None
 
 
-def safe_load_entrypoints(entrypoints: EntryPoints) -> Dict[str, Union[Exception, Any]]:
+def safe_load_entrypoints(entrypoints: EntryPoints) -> dict[str, Union[Exception, Any]]:
     """
     Load entry points for a group capturing any exceptions that occur.
     """
@@ -26,7 +26,7 @@ def safe_load_entrypoints(entrypoints: EntryPoints) -> Dict[str, Union[Exception
     #       also want to validate the type for the group for entrypoints that have
     #       a specific type we expect.
 
-    results = {}
+    results: dict[str, Union[Exception, Any]] = {}
 
     for entrypoint in entrypoints:
         result = None
@@ -40,18 +40,20 @@ def safe_load_entrypoints(entrypoints: EntryPoints) -> Dict[str, Union[Exception
     return results
 
 
-def load_prefect_collections() -> Dict[str, Union[ModuleType, Exception]]:
+def load_prefect_collections() -> dict[str, Union[ModuleType, Exception]]:
     """
     Load all Prefect collections that define an entrypoint in the group
     `prefect.collections`.
     """
-    global COLLECTIONS
+    global _collections
 
-    if COLLECTIONS is not None:
-        return COLLECTIONS
+    if _collections is not None:
+        return _collections
 
     collection_entrypoints: EntryPoints = entry_points(group="prefect.collections")
-    collections = safe_load_entrypoints(collection_entrypoints)
+    collections: dict[str, Union[Exception, Any]] = safe_load_entrypoints(
+        collection_entrypoints
+    )
 
     # TODO: Consider the utility of this once we've established this pattern.
     #       We cannot use a logger here because logging is not yet initialized.
@@ -68,5 +70,5 @@ def load_prefect_collections() -> Dict[str, Union[ModuleType, Exception]]:
             if prefect.settings.PREFECT_DEBUG_MODE:
                 print(f"Loaded collection {name!r}.")
 
-    COLLECTIONS = collections
+    _collections = collections
     return collections

prefect/results.py

@@ -245,7 +245,7 @@ async def _call_explicitly_async_block_method(
     see https://github.com/PrefectHQ/prefect/issues/15008
     """
     if hasattr(block, f"a{method}"):  # explicit async method
-        return await getattr(block.__class__.__name__, f"a{method}")(*args, **kwargs)
+        return await getattr(block, f"a{method}")(*args, **kwargs)
     elif hasattr(getattr(block, method, None), "aio"):  # sync_compatible
         return await getattr(block, method).aio(block, *args, **kwargs)
     else:  # should not happen in prefect, but users can override impls

prefect/telemetry/processors.py

@@ -1,6 +1,6 @@
 import time
 from threading import Event, Lock, Thread
-from typing import TYPE_CHECKING, Dict, Optional
+from typing import TYPE_CHECKING, Optional
 
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import Span, SpanProcessor
@@ -13,7 +13,7 @@ if TYPE_CHECKING:
 class InFlightSpanProcessor(SpanProcessor):
     def __init__(self, span_exporter: "SpanExporter"):
         self.span_exporter = span_exporter
-        self._in_flight: Dict[int, Span] = {}
+        self._in_flight: dict[int, Span] = {}
         self._lock = Lock()
         self._stop_event = Event()
         self._export_thread = Thread(target=self._export_periodically, daemon=True)
@@ -30,10 +30,10 @@ class InFlightSpanProcessor(SpanProcessor):
                     self.span_exporter.export(to_export)
 
     def _readable_span(self, span: "Span") -> "ReadableSpan":
-        readable = span._readable_span()
-        readable._end_time = time.time_ns()
-        readable._attributes = {
-            **(readable._attributes or {}),
+        readable = span._readable_span()  # pyright: ignore[reportPrivateUsage]
+        readable._end_time = time.time_ns()  # pyright: ignore[reportPrivateUsage]
+        readable._attributes = {  # pyright: ignore[reportPrivateUsage]
+            **(readable._attributes or {}),  # pyright: ignore[reportPrivateUsage]
             "prefect.in-flight": True,
         }
         return readable

prefect/telemetry/services.py

@@ -1,32 +1,38 @@
-from abc import abstractmethod
-from typing import Union
+from collections.abc import Sequence
+from typing import Any, Protocol, TypeVar
 
 from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk._logs import LogData
 from opentelemetry.sdk._logs.export import LogExporter
 from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.sdk.trace.export import SpanExporter
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 
 from prefect._internal.concurrency.services import BatchedQueueService
 
+BatchItem = TypeVar("BatchItem", ReadableSpan, LogData)
+T_contra = TypeVar("T_contra", contravariant=True)
 
-class BaseQueueingExporter(BatchedQueueService):
+
+class OTLPExporter(Protocol[T_contra]):
+    def export(self, __items: Sequence[T_contra]) -> Any:
+        ...
+
+    def shutdown(self) -> Any:
+        ...
+
+
+class BaseQueueingExporter(BatchedQueueService[BatchItem]):
     _max_batch_size = 512
     _min_interval = 2.0
-    _otlp_exporter: Union[SpanExporter, LogExporter]
-
-    def export(self, batch: list[Union[ReadableSpan, LogData]]) -> None:
-        for item in batch:
-            self.send(item)
 
-    @abstractmethod
-    def _export_batch(self, items: list[Union[ReadableSpan, LogData]]) -> None:
-        pass
+    def __init__(self, otlp_exporter: OTLPExporter[BatchItem]) -> None:
+        super().__init__()
+        self._otlp_exporter = otlp_exporter
 
-    async def _handle_batch(self, items: list[Union[ReadableSpan, LogData]]) -> None:
+    async def _handle_batch(self, items: list[BatchItem]) -> None:
         try:
-            self._export_batch(items)
+            self._otlp_exporter.export(items)
         except Exception as e:
             self._logger.exception(f"Failed to export batch: {e}")
             raise
@@ -39,29 +45,24 @@ class BaseQueueingExporter(BatchedQueueService):
         self._otlp_exporter.shutdown()
 
 
-class QueueingSpanExporter(BaseQueueingExporter, SpanExporter):
+class QueueingSpanExporter(BaseQueueingExporter[ReadableSpan], SpanExporter):
     _otlp_exporter: OTLPSpanExporter
 
     def __init__(self, endpoint: str, headers: tuple[tuple[str, str]]):
-        super().__init__()
-        self._otlp_exporter = OTLPSpanExporter(
-            endpoint=endpoint,
-            headers=dict(headers),
-        )
+        super().__init__(OTLPSpanExporter(endpoint=endpoint, headers=dict(headers)))
 
-    def _export_batch(self, items: list[ReadableSpan]) -> None:
-        self._otlp_exporter.export(items)
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        for item in spans:
+            self.send(item)
+        return SpanExportResult.SUCCESS
 
 
-class QueueingLogExporter(BaseQueueingExporter, LogExporter):
+class QueueingLogExporter(BaseQueueingExporter[LogData], LogExporter):
     _otlp_exporter: OTLPLogExporter
 
-    def __init__(self, endpoint: str, headers: tuple[tuple[str, str]]):
-        super().__init__()
-        self._otlp_exporter = OTLPLogExporter(
-            endpoint=endpoint,
-            headers=dict(headers),
-        )
+    def __init__(self, endpoint: str, headers: tuple[tuple[str, str]]) -> None:
+        super().__init__(OTLPLogExporter(endpoint=endpoint, headers=dict(headers)))
 
-    def _export_batch(self, items: list[LogData]) -> None:
-        self._otlp_exporter.export(items)
+    def export(self, batch: Sequence[LogData]) -> None:
+        for item in batch:
+            self.send(item)

prefect/workers/__init__.py

@@ -1 +1,3 @@
 from .process import ProcessWorker
+
+__all__ = ["ProcessWorker"]

{prefect_client-3.1.9.dist-info → prefect_client-3.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.1.9
+Version: 3.1.10
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.

{prefect_client-3.1.9.dist-info → prefect_client-3.1.10.dist-info}/RECORD

@@ -1,22 +1,22 @@
 prefect/.prefectignore,sha256=awSprvKT0vI8a64mEOLrMxhxqcO-b0ERQeYpA2rNKVQ,390
 prefect/__init__.py,sha256=FmdMSNpGH8Mrkn5X0mNZup8_SHdeB_aqEmS5taeOHAQ,3530
-prefect/_version.py,sha256=TTbVIJVdD1KpTcWtqBYl_bfNt7ct2xNjSPTYGZCtKGk,496
+prefect/_version.py,sha256=bUcLcvgSrFk7uLI6vWR4M14nKrDGJAXv4mOuUqQtH0U,497
 prefect/agent.py,sha256=qyyUMdiv5ndUIk-O8uwamESJGXXDQ_BmhKiLlm31ue0,286
-prefect/artifacts.py,sha256=dsxFWmdg2r9zbHM3KgKOR5YbJ29_dXUYF9kipJpbxkE,13009
-prefect/automations.py,sha256=T8sUqDriABSuFeuoKUH2OXeCK5YwFfk-mjjM0_Oflyw,5984
+prefect/artifacts.py,sha256=c_hMYTQoKFij8_tahtBeHzjYl7ztXAU90QP_6SuMw9A,12934
+prefect/automations.py,sha256=QHS3Xh6wUKRvzxXt8cfWri6wXjX8EfEDDQN8a7cRw3c,6594
 prefect/cache_policies.py,sha256=dlhYHMxhb2N5KdDMIAuw4q-PIsI8WH5k0LDDLqsRYMU,9903
-prefect/context.py,sha256=HTEqJIvd_VuSS9LYdhZ2qlZsPgftWndgkw6EFGKbWJw,22364
+prefect/context.py,sha256=OAEhyJz7kheIh6UEwPp3oLklYdM_UHibrdqeM3aCLqk,22325
 prefect/engine.py,sha256=qkT6hQWfxQPmjAEECQx3wluiDiFMHf3h9DRG1ShwK7w,2031
 prefect/exceptions.py,sha256=sbphPKQ4yOBUa9w0MsSFoDj_uC8Tlv9WHTjzO3cQKq8,11593
 prefect/filesystems.py,sha256=cLBGbnW2NQqH4N70Y1i1QuCL3R7p-_r_d5CGaEAd4wg,17897
 prefect/flow_engine.py,sha256=eb2N5lRdsZBWNIsaAs1pqdDhL5Vevq8ETxyn5LM4Mks,53812
 prefect/flow_runs.py,sha256=-5udBBYdgdCBCjAMYvELbA1vmrjZ6oREXl-BZdZr6hc,16129
-prefect/flows.py,sha256=FaArZfFYQgPnKRCC4c3nsEfJlwQTBQrCM7UHhULmrl4,94110
-prefect/futures.py,sha256=DlZvdccKtwQKuDUFrZ4zcINeO9C1chLiNOwjE5gTgCk,16356
+prefect/flows.py,sha256=Z1BdpmJotXJoqoDqcUtxiiPsGWvDfSBvk1XEEo6ykw4,94127
+prefect/futures.py,sha256=ZtQiRJddO5-LNabUZiX_2wEJnRf6w7qZ8KRC5VUNOa0,16915
 prefect/main.py,sha256=lFYvXkW32sMj4aVH97foApD3Ut3rZIwtO91xUQHo_6I,2355
-prefect/plugins.py,sha256=HY7Z7OJlltqzsUiPMEL1Y_hQbHw0CeZKayWiK-k8DP4,2435
+prefect/plugins.py,sha256=FPRLR2mWVBMuOnlzeiTD9krlHONZH2rtYLD753JQDNQ,2516
 prefect/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-prefect/results.py,sha256=TiGitDQmG4wnrnxhV04mKnSe30gczGstLQa6ItrAOyM,50581
+prefect/results.py,sha256=okLkMlQcYUZHIl0XXMqVnWVLdcIYXe4b1wWX1T7Qa_8,50562
 prefect/serializers.py,sha256=DiKLYdIA5VjCSrLyu0BSWcLLVRujqjvj8XoQ3HOuvoY,8968
 prefect/states.py,sha256=ZhJHLx8ax7dQNdaBl2eDT6VYGAU-68R1loht68qHGpc,25396
 prefect/task_engine.py,sha256=1xM64s9OpDREnKPnU9HpwHgtbAE50uwdp2ZucmPUhCg,61226
@@ -133,9 +133,9 @@ prefect/input/__init__.py,sha256=Ue2h-YhYP71nEtsVJaslqMwO6C0ckjhjTYwwEgp-E3g,701
 prefect/input/actions.py,sha256=IGdWjVcesnRjLmPCzB4ZM7FkRWXDKCku6yhE-7p0vKk,3777
 prefect/input/run_input.py,sha256=2wG-0L3N0spwh61Z3xI0PM8AAjHEIQZcDN703Er_gLo,18728
 prefect/locking/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-prefect/locking/filesystem.py,sha256=GiZlLLj51cLH6QQgq7IeU6jUK6vGi0wMnOG0zaO95-c,8025
-prefect/locking/memory.py,sha256=Y1fsMSUAk3jUILzRivbxlrE9Xv8OcVbaylVf-aiEGNc,7495
-prefect/locking/protocol.py,sha256=o5-48SxvEDAdVwW8RIn7rCN32CmvIsaVHTztESUXuHU,4232
+prefect/locking/filesystem.py,sha256=zhNwdKroi2kLR6Cut6CMT-rWmFwtTtzuGKSwGH_Iw0s,8084
+prefect/locking/memory.py,sha256=mFUgV750ywEL7aVQuxFjg9gxbjVU4esBQn7bGQYzeMY,7548
+prefect/locking/protocol.py,sha256=RsfvlaHTTEJ0YvYWSqFGoZuT2w4FPPxyQlHqjoyNGuE,4240
 prefect/logging/__init__.py,sha256=zx9f5_dWrR4DbcTOFBpNGOPoCZ1QcPFudr7zxb2XRpA,148
 prefect/logging/configuration.py,sha256=t7EJrk7Jr4QMxopH3TP6gRN8kMg_LFBsIoiwXCihrso,3353
 prefect/logging/filters.py,sha256=9keHLN4-cpnsWcii1qU0RITNi9-m7pOhkJ_t0MtCM4k,1117
@@ -200,9 +200,9 @@ prefect/telemetry/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs
 prefect/telemetry/bootstrap.py,sha256=b64UpBBo0Yo7Z8b_V9NQLAlFrpIuJoTs5opYTuB-R20,1550
 prefect/telemetry/instrumentation.py,sha256=G8kDiN2N4l7A8on-BHbYnh8uqH71zuGr3t2JntxRv0Q,4183
 prefect/telemetry/logging.py,sha256=yn5D4D2GGRrAv0y8wlHPN7PZDmQucGjQT_YauK9M9Yo,727
-prefect/telemetry/processors.py,sha256=mBWk51uIcdFumus0L_FZ0Qf-0c43jzDUAkoNETH_wJ8,2122
+prefect/telemetry/processors.py,sha256=jw6j6LviOVxw3IBJe7cSjsxFk0zzY43jUmy6C9pcfCE,2272
 prefect/telemetry/run_telemetry.py,sha256=FfvcUJK6Sqcagp6rJ3vx4twCAkyYJYBhtNUUzjM0D7A,8135
-prefect/telemetry/services.py,sha256=eGKcg9pKUsTnrs9ou-vOCn1wXN2f-0iD9b1Upk01R3o,2146
+prefect/telemetry/services.py,sha256=ek5KMSgCjUeIoNl9INPCfKkwlUwyFYp_dUiZd_rD0_8,2270
 prefect/types/__init__.py,sha256=s-l0r8UiDs7GCEstOPMzbb9PvvF_0F4G8OcVy3_50Yk,4700
 prefect/types/entrypoint.py,sha256=2FF03-wLPgtnqR_bKJDB2BsXXINPdu8ptY9ZYEZnXg8,328
 prefect/utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -234,15 +234,15 @@ prefect/utilities/visualization.py,sha256=4MyQKLb00A1QZCvSrbPEbyqUNByQ_EkJ6wXKFC
 prefect/utilities/schema_tools/__init__.py,sha256=At3rMHd2g_Em2P3_dFQlFgqR_EpBwrYtU2N2OJd0pDE,345
 prefect/utilities/schema_tools/hydration.py,sha256=4yQMynZSloPFp0tlA7g8udWqE2TjZgfm19Y4R4F0s04,9415
 prefect/utilities/schema_tools/validation.py,sha256=Wix26IVR-ZJ32-6MX2pHhrwm3reB-Q4iB6_phn85OKE,10743
-prefect/workers/__init__.py,sha256=8dP8SLZbWYyC_l9DRTQSE3dEbDgns5DZDhxkp_NfsbQ,35
+prefect/workers/__init__.py,sha256=EaM1F0RZ-XIJaGeTKLsXDnfOPHzVWk5bk0_c4BVS44M,64
 prefect/workers/base.py,sha256=nBv4pc51J1JRoSkgxvBDfQb4wIEN67D7cBVwT-jYHdI,49198
 prefect/workers/block.py,sha256=qyyUMdiv5ndUIk-O8uwamESJGXXDQ_BmhKiLlm31ue0,286
 prefect/workers/cloud.py,sha256=qyyUMdiv5ndUIk-O8uwamESJGXXDQ_BmhKiLlm31ue0,286
 prefect/workers/process.py,sha256=tcJ3fbiraLCfpVGpv8dOHwMSfVzeD_kyguUOvPuIz6I,19796
 prefect/workers/server.py,sha256=lgh2FfSuaNU7b6HPxSFm8JtKvAvHsZGkiOo4y4tW1Cw,2022
 prefect/workers/utilities.py,sha256=VfPfAlGtTuDj0-Kb8WlMgAuOfgXCdrGAnKMapPSBrwc,2483
-prefect_client-3.1.9.dist-info/LICENSE,sha256=MCxsn8osAkzfxKC4CC_dLcUkU8DZLkyihZ8mGs3Ah3Q,11357
-prefect_client-3.1.9.dist-info/METADATA,sha256=KmvB9sA6WKsVD23CjmL-BLNkORJIbhyKv28iSnvbrxE,7286
-prefect_client-3.1.9.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-prefect_client-3.1.9.dist-info/top_level.txt,sha256=MJZYJgFdbRc2woQCeB4vM6T33tr01TmkEhRcns6H_H4,8
-prefect_client-3.1.9.dist-info/RECORD,,
+prefect_client-3.1.10.dist-info/LICENSE,sha256=MCxsn8osAkzfxKC4CC_dLcUkU8DZLkyihZ8mGs3Ah3Q,11357
+prefect_client-3.1.10.dist-info/METADATA,sha256=BbPdwx1OPrG1CxVh7OJNeBzquSMMMfOCBk4e-zilqhs,7287
+prefect_client-3.1.10.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+prefect_client-3.1.10.dist-info/top_level.txt,sha256=MJZYJgFdbRc2woQCeB4vM6T33tr01TmkEhRcns6H_H4,8
+prefect_client-3.1.10.dist-info/RECORD,,