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

prefect/__init__.py

@@ -2,29 +2,14 @@
 
 # Setup version and path constants
 
+import sys
 from . import _version
 import importlib
 import pathlib
-from typing import TYPE_CHECKING, Any
-
-__version_info__ = _version.get_versions()
-__version__ = __version_info__["version"]
-
-# The absolute path to this module
-__module_path__ = pathlib.Path(__file__).parent
-# The absolute path to the root of the repository, only valid for use during development
-__development_base_path__ = __module_path__.parents[1]
-
-# The absolute path to the built UI within the Python module, used by
-# `prefect server start` to serve a dynamic build of the UI
-__ui_static_subpath__ = __module_path__ / "server" / "ui_build"
-
-# The absolute path to the built UI within the Python module
-__ui_static_path__ = __module_path__ / "server" / "ui"
-
-del _version, pathlib
+from typing import TYPE_CHECKING, Any, Optional, TypedDict, cast
 
 if TYPE_CHECKING:
+    from importlib.machinery import ModuleSpec
     from .main import (
         allow_failure,
         flow,
@@ -45,80 +30,89 @@ if TYPE_CHECKING:
         suspend_flow_run,
     )
 
-_slots: dict[str, Any] = {
-    "__version_info__": __version_info__,
-    "__version__": __version__,
-    "__module_path__": __module_path__,
-    "__development_base_path__": __development_base_path__,
-    "__ui_static_subpath__": __ui_static_subpath__,
-    "__ui_static_path__": __ui_static_path__,
-}
+    __spec__: ModuleSpec
+
+    # Versioneer provides version information as dictionaries
+    # with these keys
+    class VersionInfo(TypedDict("_FullRevisionId", {"full-revisionid": str})):
+        version: str
+        dirty: Optional[bool]
+        error: Optional[str]
+        date: Optional[str]
 
-_public_api: dict[str, tuple[str, str]] = {
+
+__version_info__: "VersionInfo" = cast("VersionInfo", _version.get_versions())
+__version__ = __version_info__["version"]
+
+# The absolute path to this module
+__module_path__: pathlib.Path = pathlib.Path(__file__).parent
+# The absolute path to the root of the repository, only valid for use during development
+__development_base_path__: pathlib.Path = __module_path__.parents[1]
+
+# The absolute path to the built UI within the Python module, used by
+# `prefect server start` to serve a dynamic build of the UI
+__ui_static_subpath__: pathlib.Path = __module_path__ / "server" / "ui_build"
+
+# The absolute path to the built UI within the Python module
+__ui_static_path__: pathlib.Path = __module_path__ / "server" / "ui"
+
+del _version, pathlib
+
+_public_api: dict[str, tuple[Optional[str], str]] = {
     "allow_failure": (__spec__.parent, ".main"),
+    "aserve": (__spec__.parent, ".main"),
+    "deploy": (__spec__.parent, ".main"),
     "flow": (__spec__.parent, ".main"),
     "Flow": (__spec__.parent, ".main"),
     "get_client": (__spec__.parent, ".main"),
     "get_run_logger": (__spec__.parent, ".main"),
+    "pause_flow_run": (__spec__.parent, ".main"),
+    "resume_flow_run": (__spec__.parent, ".main"),
+    "serve": (__spec__.parent, ".main"),
     "State": (__spec__.parent, ".main"),
+    "suspend_flow_run": (__spec__.parent, ".main"),
     "tags": (__spec__.parent, ".main"),
     "task": (__spec__.parent, ".main"),
     "Task": (__spec__.parent, ".main"),
     "Transaction": (__spec__.parent, ".main"),
     "unmapped": (__spec__.parent, ".main"),
-    "serve": (__spec__.parent, ".main"),
-    "aserve": (__spec__.parent, ".main"),
-    "deploy": (__spec__.parent, ".main"),
-    "pause_flow_run": (__spec__.parent, ".main"),
-    "resume_flow_run": (__spec__.parent, ".main"),
-    "suspend_flow_run": (__spec__.parent, ".main"),
 }
 
 # Declare API for type-checkers
 __all__ = [
+    "__version__",
     "allow_failure",
+    "aserve",
+    "deploy",
     "flow",
     "Flow",
     "get_client",
     "get_run_logger",
+    "pause_flow_run",
+    "resume_flow_run",
+    "serve",
     "State",
+    "suspend_flow_run",
     "tags",
     "task",
     "Task",
     "Transaction",
     "unmapped",
-    "serve",
-    "aserve",
-    "deploy",
-    "pause_flow_run",
-    "resume_flow_run",
-    "suspend_flow_run",
-    "__version_info__",
-    "__version__",
-    "__module_path__",
-    "__development_base_path__",
-    "__ui_static_subpath__",
-    "__ui_static_path__",
 ]
 
 
-def __getattr__(attr_name: str) -> object:
-    if attr_name in _slots:
-        return _slots[attr_name]
+def __getattr__(attr_name: str) -> Any:
     try:
-        dynamic_attr = _public_api.get(attr_name)
-        if dynamic_attr is None:
+        if (dynamic_attr := _public_api.get(attr_name)) is None:
             return importlib.import_module(f".{attr_name}", package=__name__)
 
-        package, module_name = dynamic_attr
-
-        from importlib import import_module
-
-        if module_name == "__module__":
-            return import_module(f".{attr_name}", package=package)
+        package, mname = dynamic_attr
+        if mname == "__module__":
+            return importlib.import_module(f".{attr_name}", package=package)
         else:
-            module = import_module(module_name, package=package)
+            module = importlib.import_module(mname, package=package)
             return getattr(module, attr_name)
     except ModuleNotFoundError as ex:
-        module, _, attribute = ex.name.rpartition(".")
-        raise AttributeError(f"module {module} has no attribute {attribute}") from ex
+        mname, _, attr = (ex.name or "").rpartition(".")
+        ctx = {"name": mname, "obj": attr} if sys.version_info >= (3, 10) else {}
+        raise AttributeError(f"module {mname} has no attribute {attr}", **ctx) from ex

prefect/_internal/concurrency/services.py

@@ -225,7 +225,9 @@ class QueueService(abc.ABC, Generic[T]):
             return future.result()
 
     @classmethod
-    def drain_all(cls, timeout: Optional[float] = None) -> Union[Awaitable, None]:
+    def drain_all(
+        cls, timeout: Optional[float] = None, at_exit=True
+    ) -> Union[Awaitable, None]:
         """
         Stop all instances of the service and wait for all remaining work to be
         completed.
@@ -237,7 +239,7 @@ class QueueService(abc.ABC, Generic[T]):
             instances = tuple(cls._instances.values())
 
             for instance in instances:
-                futures.append(instance._drain())
+                futures.append(instance._drain(at_exit=at_exit))
 
         if get_running_loop() is not None:
             return (
@@ -376,10 +378,10 @@ class BatchedQueueService(QueueService[T]):
 @contextlib.contextmanager
 def drain_on_exit(service: QueueService):
     yield
-    service.drain_all()
+    service.drain_all(at_exit=True)
 
 
 @contextlib.asynccontextmanager
 async def drain_on_exit_async(service: QueueService):
     yield
-    await service.drain_all()
+    await service.drain_all(at_exit=True)

prefect/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-17T10:20:43-0800",
+ "date": "2024-12-24T22:58:40-0800",
  "dirty": true,
  "error": null,
- "full-revisionid": "53a83ebc4a9a26a1d7cfd91d6be7cf6bd30e7f21",
- "version": "3.1.8"
+ "full-revisionid": "b11b56fc6143df2bd4c7018f7a974dc656fe9aa3",
+ "version": "3.1.10"
 }
 '''  # END VERSION_JSON
 

prefect/agent.py

@@ -1,6 +1,8 @@
 """
 2024-06-27: This surfaces an actionable error message for moved or removed objects in Prefect 3.0 upgrade.
 """
+from typing import Any, Callable
+
 from prefect._internal.compatibility.migration import getattr_migration
 
-__getattr__ = getattr_migration(__name__)
+__getattr__: Callable[[str], Any] = getattr_migration(__name__)

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/client/cloud.py

@@ -1,6 +1,5 @@
 import re
 from typing import Any, NoReturn, Optional, cast
-from uuid import UUID
 
 import anyio
 import httpx
@@ -23,7 +22,6 @@ from prefect.settings import (
     PREFECT_CLOUD_API_URL,
     PREFECT_TESTING_UNIT_TEST_MODE,
 )
-from prefect.types import KeyValueLabels
 
 PARSE_API_URL_REGEX = re.compile(r"accounts/(.{36})/workspaces/(.{36})")
 
@@ -160,25 +158,6 @@ class CloudClient:
         response = await self.get(f"{self.account_base_url}/ip_allowlist/my_access")
         return IPAllowlistMyAccessResponse.model_validate(response)
 
-    async def update_flow_run_labels(
-        self, flow_run_id: UUID, labels: KeyValueLabels
-    ) -> httpx.Response:
-        """
-        Update the labels for a flow run.
-
-        Args:
-            flow_run_id: The identifier for the flow run to update.
-            labels: A dictionary of labels to update for the flow run.
-
-        Returns:
-            an `httpx.Response` object from the PATCH request
-        """
-
-        return await self._client.patch(
-            f"{self.workspace_base_url}/flow_runs/{flow_run_id}/labels",
-            json=labels,
-        )
-
     async def __aenter__(self) -> Self:
         await self._client.__aenter__()
         return self

prefect/client/schemas/objects.py

@@ -190,6 +190,8 @@ class StateDetails(PrefectBaseModel):
     retriable: Optional[bool] = None
     transition_id: Optional[UUID] = None
     task_parameters_id: Optional[UUID] = None
+    # Captures the trace_id and span_id of the span where this state was created
+    traceparent: Optional[str] = None
 
 
 def data_discriminator(x: Any) -> str:
@@ -237,6 +239,15 @@ class State(ObjectBaseModel, Generic[R]):
     ) -> Union[R, Exception]:
         ...
 
+    @overload
+    def result(
+        self: "State[R]",
+        raise_on_failure: bool = ...,
+        fetch: bool = ...,
+        retry_result_failure: bool = ...,
+    ) -> Union[R, Exception]:
+        ...
+
     @deprecated.deprecated_parameter(
         "fetch",
         when=lambda fetch: fetch is not True,

prefect/client/utilities.py

@@ -7,7 +7,7 @@ Utilities for working with clients.
 
 from collections.abc import Awaitable, Coroutine
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
 from typing_extensions import Concatenate, ParamSpec, TypeGuard, TypeVar
 
@@ -71,23 +71,9 @@ def client_injector(
     return wrapper
 
 
-@overload
 def inject_client(
     fn: Callable[P, Coroutine[Any, Any, R]],
 ) -> Callable[P, Coroutine[Any, Any, R]]:
-    ...
-
-
-@overload
-def inject_client(
-    fn: Callable[P, R],
-) -> Callable[P, R]:
-    ...
-
-
-def inject_client(
-    fn: Callable[P, Union[Coroutine[Any, Any, R], R]],
-) -> Callable[P, Union[Coroutine[Any, Any, R], R]]:
     """
     Simple helper to provide a context managed client to an asynchronous function.