prefect-client 3.1.4__py3-none-any.whl → 3.1.6__py3-none-any.whl

prefect/_internal/schemas/bases.py

@@ -4,7 +4,7 @@ Utilities for creating and working with Prefect REST API schemas.
 
 import datetime
 import os
-from typing import Any, ClassVar, Optional, Set, TypeVar
+from typing import Any, ClassVar, Generator, Optional, Set, TypeVar, cast
 from uuid import UUID, uuid4
 
 import pendulum
@@ -33,7 +33,7 @@ class PrefectBaseModel(BaseModel):
 
     _reset_fields: ClassVar[Set[str]] = set()
 
-    model_config = ConfigDict(
+    model_config: ClassVar[ConfigDict] = ConfigDict(
         ser_json_timedelta="float",
         defer_build=True,
         extra=(
@@ -58,7 +58,7 @@ class PrefectBaseModel(BaseModel):
         else:
             return copy_dict == other
 
-    def __rich_repr__(self):
+    def __rich_repr__(self) -> Generator[tuple[str, Any, Any], None, None]:
         # Display all of the fields in the model if they differ from the default value
         for name, field in self.model_fields.items():
             value = getattr(self, name)
@@ -71,9 +71,11 @@ class PrefectBaseModel(BaseModel):
                 and name == "timestamp"
                 and value
             ):
-                value = pendulum.instance(value).isoformat()
+                value = cast(pendulum.DateTime, pendulum.instance(value)).isoformat()
             elif isinstance(field.annotation, datetime.datetime) and value:
-                value = pendulum.instance(value).diff_for_humans()
+                value = cast(
+                    pendulum.DateTime, pendulum.instance(value)
+                ).diff_for_humans()
 
             yield name, value, field.get_default()
 
@@ -84,12 +86,9 @@ class PrefectBaseModel(BaseModel):
         Returns:
             PrefectBaseModel: A new instance of the model with the reset fields.
         """
-        data = self.model_dump()
         return self.model_copy(
             update={
-                field: self.model_fields[field].get_default(
-                    call_default_factory=True, validated_data=data
-                )
+                field: self.model_fields[field].get_default(call_default_factory=True)
                 for field in self._reset_fields
             }
         )
@@ -116,11 +115,11 @@ class ObjectBaseModel(IDBaseModel):
     """
 
     _reset_fields: ClassVar[Set[str]] = {"id", "created", "updated"}
-    model_config = ConfigDict(from_attributes=True)
+    model_config: ClassVar[ConfigDict] = ConfigDict(from_attributes=True)
 
     created: Optional[DateTime] = Field(default=None, repr=False)
     updated: Optional[DateTime] = Field(default=None, repr=False)
 
 
 class ActionBaseModel(PrefectBaseModel):
-    model_config: ConfigDict = ConfigDict(extra="forbid")
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")

prefect/_internal/schemas/validators.py

@@ -13,6 +13,7 @@ import warnings
 from copy import copy
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union
+from uuid import UUID
 
 import jsonschema
 import pendulum
@@ -653,7 +654,7 @@ def validate_message_template_variables(v: Optional[str]) -> Optional[str]:
     return v
 
 
-def validate_default_queue_id_not_none(v: Optional[str]) -> Optional[str]:
+def validate_default_queue_id_not_none(v: Optional[UUID]) -> UUID:
     if v is None:
         raise ValueError(
             "`default_queue_id` is a required field. If you are "

prefect/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-11-20T19:37:35-0600",
+ "date": "2024-12-11T10:46:23-0500",
  "dirty": true,
  "error": null,
- "full-revisionid": "78ee41cb26058cc2dfba8aece6a2471ec5ea3675",
- "version": "3.1.4"
+ "full-revisionid": "6bfce9e87c6e7e3e44c65d8b5ecc6ab156f2c7c8",
+ "version": "3.1.6"
 }
 '''  # END VERSION_JSON
 

prefect/automations.py

@@ -1,10 +1,10 @@
-from typing import Optional
+from typing import Optional, Type
 from uuid import UUID
 
 from pydantic import Field
 from typing_extensions import Self
 
-from prefect.client.utilities import get_or_create_client
+from prefect.client.orchestration import get_client
 from prefect.events.actions import (
     CallWebhook,
     CancelFlowRun,
@@ -99,10 +99,10 @@ class Automation(AutomationCore):
         )
         created_automation = auto_to_create.create()
         """
-        client, _ = get_or_create_client()
-        automation = AutomationCore(**self.model_dump(exclude={"id"}))
-        self.id = await client.create_automation(automation=automation)
-        return self
+        async with get_client() as client:
+            automation = AutomationCore(**self.model_dump(exclude={"id"}))
+            self.id = await client.create_automation(automation=automation)
+            return self
 
     @sync_compatible
     async def update(self: Self):
@@ -112,15 +112,16 @@ class Automation(AutomationCore):
         auto.name = "new name"
         auto.update()
         """
-
-        client, _ = get_or_create_client()
-        automation = AutomationCore(**self.model_dump(exclude={"id", "owner_resource"}))
-        await client.update_automation(automation_id=self.id, automation=automation)
+        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)
 
     @classmethod
     @sync_compatible
     async def read(
-        cls: Self, id: Optional[UUID] = None, name: Optional[str] = None
+        cls: Type[Self], id: Optional[UUID] = None, name: Optional[str] = None
     ) -> Self:
         """
         Read an automation by ID or name.
@@ -134,20 +135,25 @@ class Automation(AutomationCore):
             raise ValueError("Only one of id or name can be provided")
         if not id and not name:
             raise ValueError("One of id or name must be provided")
-        client, _ = get_or_create_client()
-        if id:
-            try:
-                automation = await client.read_automation(automation_id=id)
-            except PrefectHTTPStatusError as exc:
-                if exc.response.status_code == 404:
+        async with get_client() as client:
+            if id:
+                try:
+                    automation = await client.read_automation(automation_id=id)
+                except PrefectHTTPStatusError as exc:
+                    if exc.response.status_code == 404:
+                        raise ValueError(f"Automation with ID {id!r} not found")
+                    raise
+                if automation is None:
                     raise ValueError(f"Automation with ID {id!r} not found")
-            return Automation(**automation.model_dump())
-        else:
-            automation = await client.read_automations_by_name(name=name)
-            if len(automation) > 0:
-                return Automation(**automation[0].model_dump()) if automation else None
+                return Automation(**automation.model_dump())
             else:
-                raise ValueError(f"Automation with name {name!r} not found")
+                automation = await client.read_automations_by_name(name=name)
+                if len(automation) > 0:
+                    return (
+                        Automation(**automation[0].model_dump()) if automation else None
+                    )
+                else:
+                    raise ValueError(f"Automation with name {name!r} not found")
 
     @sync_compatible
     async def delete(self: Self) -> bool:
@@ -155,14 +161,14 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.delete()
         """
-        try:
-            client, _ = get_or_create_client()
-            await client.delete_automation(self.id)
-            return True
-        except PrefectHTTPStatusError as exc:
-            if exc.response.status_code == 404:
-                return False
-            raise
+        async with get_client() as client:
+            try:
+                await client.delete_automation(self.id)
+                return True
+            except PrefectHTTPStatusError as exc:
+                if exc.response.status_code == 404:
+                    return False
+                raise
 
     @sync_compatible
     async def disable(self: Self) -> bool:
@@ -171,14 +177,14 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.disable()
         """
-        try:
-            client, _ = get_or_create_client()
-            await client.pause_automation(self.id)
-            return True
-        except PrefectHTTPStatusError as exc:
-            if exc.response.status_code == 404:
-                return False
-            raise
+        async with get_client() as client:
+            try:
+                await client.pause_automation(self.id)
+                return True
+            except PrefectHTTPStatusError as exc:
+                if exc.response.status_code == 404:
+                    return False
+                raise
 
     @sync_compatible
     async def enable(self: Self) -> bool:
@@ -187,11 +193,11 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.enable()
         """
-        try:
-            client, _ = get_or_create_client()
-            await client.resume_automation("asd")
-            return True
-        except PrefectHTTPStatusError as exc:
-            if exc.response.status_code == 404:
-                return False
-            raise
+        async with get_client() as client:
+            try:
+                await client.resume_automation(self.id)
+                return True
+            except PrefectHTTPStatusError as exc:
+                if exc.response.status_code == 404:
+                    return False
+                raise

prefect/blocks/abstract.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from logging import Logger
+from logging import Logger, LoggerAdapter
 from pathlib import Path
 from typing import (
     Any,
@@ -15,7 +15,7 @@ from typing import (
     Union,
 )
 
-from typing_extensions import Self
+from typing_extensions import Self, TypeAlias
 
 from prefect.blocks.core import Block
 from prefect.exceptions import MissingContextError
@@ -23,6 +23,8 @@ from prefect.logging.loggers import get_logger, get_run_logger
 
 T = TypeVar("T")
 
+LoggerOrAdapter: TypeAlias = Union[Logger, LoggerAdapter]
+
 
 class CredentialsBlock(Block, ABC):
     """
@@ -34,7 +36,7 @@ class CredentialsBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the CredentialsBlock
         is called from within a flow or task run context.
@@ -73,10 +75,10 @@ class NotificationBlock(Block, ABC):
     """
 
     _block_schema_capabilities = ["notify"]
-    _events_excluded_methods = Block._events_excluded_methods.default + ["notify"]
+    _events_excluded_methods = Block._events_excluded_methods + ["notify"]
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the NotificationBlock
         is called from within a flow or task run context.
@@ -123,7 +125,7 @@ class JobRun(ABC, Generic[T]):  # not a block
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the JobRun
         is called from within a flow or task run context.
@@ -158,7 +160,7 @@ class JobBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the JobBlock
         is called from within a flow or task run context.
@@ -202,7 +204,7 @@ class DatabaseBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the DatabaseBlock
         is called from within a flow or task run context.
@@ -337,7 +339,7 @@ class ObjectStorageBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the ObjectStorageBlock
         is called from within a flow or task run context.
@@ -469,7 +471,7 @@ class SecretBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the SecretBlock
         is called from within a flow or task run context.

prefect/blocks/core.py

@@ -326,7 +326,9 @@ class Block(BaseModel, ABC):
 
     # Exclude `save` as it uses the `sync_compatible` decorator and needs to be
     # decorated directly.
-    _events_excluded_methods = ["block_initialization", "save", "dict"]
+    _events_excluded_methods: ClassVar[List[str]] = PrivateAttr(
+        default=["block_initialization", "save", "dict"]
+    )
 
     @classmethod
     def __dispatch_key__(cls):
@@ -1229,7 +1231,7 @@ class Block(BaseModel, ABC):
         """
         block_type_slug = kwargs.pop("block_type_slug", None)
         if block_type_slug:
-            subcls = lookup_type(cls, dispatch_key=block_type_slug)
+            subcls = cls.get_block_class_from_key(block_type_slug)
             return super().__new__(subcls)
         else:
             return super().__new__(cls)

prefect/cache_policies.py

@@ -75,12 +75,12 @@ class CachePolicy:
         task_ctx: TaskRunContext,
         inputs: Dict[str, Any],
         flow_parameters: Dict[str, Any],
-        **kwargs,
+        **kwargs: Any,
     ) -> Optional[str]:
         raise NotImplementedError
 
     def __sub__(self, other: str) -> "CachePolicy":
-        if not isinstance(other, str):
+        if not isinstance(other, str):  # type: ignore[reportUnnecessaryIsInstance]
             raise TypeError("Can only subtract strings from key policies.")
         new = Inputs(exclude=[other])
         return CompoundCachePolicy(policies=[self, new])
@@ -140,7 +140,7 @@ class CacheKeyFnPolicy(CachePolicy):
         task_ctx: TaskRunContext,
         inputs: Dict[str, Any],
         flow_parameters: Dict[str, Any],
-        **kwargs,
+        **kwargs: Any,
     ) -> Optional[str]:
         if self.cache_key_fn:
             return self.cache_key_fn(task_ctx, inputs)
@@ -162,9 +162,9 @@ class CompoundCachePolicy(CachePolicy):
         task_ctx: TaskRunContext,
         inputs: Dict[str, Any],
         flow_parameters: Dict[str, Any],
-        **kwargs,
+        **kwargs: Any,
     ) -> Optional[str]:
-        keys = []
+        keys: list[str] = []
         for policy in self.policies:
             policy_key = policy.compute_key(
                 task_ctx=task_ctx,
@@ -191,7 +191,7 @@ class _None(CachePolicy):
         task_ctx: TaskRunContext,
         inputs: Dict[str, Any],
         flow_parameters: Dict[str, Any],
-        **kwargs,
+        **kwargs: Any,
     ) -> Optional[str]:
         return None
 
@@ -211,7 +211,7 @@ class TaskSource(CachePolicy):
         task_ctx: TaskRunContext,
         inputs: Optional[Dict[str, Any]],
         flow_parameters: Optional[Dict[str, Any]],
-        **kwargs,
+        **kwargs: Any,
     ) -> Optional[str]:
         if not task_ctx:
             return None
@@ -238,7 +238,7 @@ class FlowParameters(CachePolicy):
         task_ctx: TaskRunContext,
         inputs: Dict[str, Any],
         flow_parameters: Dict[str, Any],
-        **kwargs,
+        **kwargs: Any,
     ) -> Optional[str]:
         if not flow_parameters:
             return None
@@ -257,7 +257,7 @@ class RunId(CachePolicy):
         task_ctx: TaskRunContext,
         inputs: Dict[str, Any],
         flow_parameters: Dict[str, Any],
-        **kwargs,
+        **kwargs: Any,
     ) -> Optional[str]:
         if not task_ctx:
             return None
@@ -280,7 +280,7 @@ class Inputs(CachePolicy):
         task_ctx: TaskRunContext,
         inputs: Dict[str, Any],
         flow_parameters: Dict[str, Any],
-        **kwargs,
+        **kwargs: Any,
     ) -> Optional[str]:
         hashed_inputs = {}
         inputs = inputs or {}
@@ -307,7 +307,7 @@ class Inputs(CachePolicy):
             raise ValueError(msg) from exc
 
     def __sub__(self, other: str) -> "CachePolicy":
-        if not isinstance(other, str):
+        if not isinstance(other, str):  # type: ignore[reportUnnecessaryIsInstance]
             raise TypeError("Can only subtract strings from key policies.")
         return Inputs(exclude=self.exclude + [other])
 

prefect/client/__init__.py

@@ -16,6 +16,8 @@ $ python -m asyncio
 </div>
 """
 
+from collections.abc import Callable
+from typing import Any
 from prefect._internal.compatibility.migration import getattr_migration
 
-__getattr__ = getattr_migration(__name__)
+__getattr__: Callable[[str], Any] = getattr_migration(__name__)

prefect/client/base.py

@@ -4,22 +4,11 @@ import threading
 import time
 import uuid
 from collections import defaultdict
+from collections.abc import AsyncGenerator, Awaitable, MutableMapping
 from contextlib import asynccontextmanager
 from datetime import datetime, timezone
-from typing import (
-    Any,
-    AsyncGenerator,
-    Awaitable,
-    Callable,
-    Dict,
-    MutableMapping,
-    Optional,
-    Protocol,
-    Set,
-    Tuple,
-    Type,
-    runtime_checkable,
-)
+from logging import Logger
+from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, runtime_checkable
 
 import anyio
 import httpx
@@ -46,14 +35,14 @@ from prefect.utilities.math import bounded_poisson_interval, clamped_poisson_int
 
 # Datastores for lifespan management, keys should be a tuple of thread and app
 # identities.
-APP_LIFESPANS: Dict[Tuple[int, int], LifespanManager] = {}
-APP_LIFESPANS_REF_COUNTS: Dict[Tuple[int, int], int] = {}
+APP_LIFESPANS: dict[tuple[int, int], LifespanManager] = {}
+APP_LIFESPANS_REF_COUNTS: dict[tuple[int, int], int] = {}
 # Blocks concurrent access to the above dicts per thread. The index should be the thread
 # identity.
-APP_LIFESPANS_LOCKS: Dict[int, anyio.Lock] = defaultdict(anyio.Lock)
+APP_LIFESPANS_LOCKS: dict[int, anyio.Lock] = defaultdict(anyio.Lock)
 
 
-logger = get_logger("client")
+logger: Logger = get_logger("client")
 
 
 # Define ASGI application types for type checking
@@ -161,7 +150,7 @@ class PrefectResponse(httpx.Response):
     Provides more informative error messages.
     """
 
-    def raise_for_status(self) -> None:
+    def raise_for_status(self) -> Response:
         """
         Raise an exception if the response contains an HTTPStatusError.
 
@@ -174,9 +163,9 @@ class PrefectResponse(httpx.Response):
             raise PrefectHTTPStatusError.from_httpx_error(exc) from exc.__cause__
 
     @classmethod
-    def from_httpx_response(cls: Type[Self], response: httpx.Response) -> Self:
+    def from_httpx_response(cls: type[Self], response: httpx.Response) -> Response:
         """
-        Create a `PrefectReponse` from an `httpx.Response`.
+        Create a `PrefectResponse` from an `httpx.Response`.
 
         By changing the `__class__` attribute of the Response, we change the method
         resolution order to look for methods defined in PrefectResponse, while leaving
@@ -200,10 +189,10 @@ class PrefectHttpxAsyncClient(httpx.AsyncClient):
 
     def __init__(
         self,
-        *args,
+        *args: Any,
         enable_csrf_support: bool = False,
         raise_on_all_errors: bool = True,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.enable_csrf_support: bool = enable_csrf_support
         self.csrf_token: Optional[str] = None
@@ -222,10 +211,10 @@ class PrefectHttpxAsyncClient(httpx.AsyncClient):
         self,
         request: Request,
         send: Callable[[Request], Awaitable[Response]],
-        send_args: Tuple,
-        send_kwargs: Dict,
-        retry_codes: Set[int] = set(),
-        retry_exceptions: Tuple[Exception, ...] = tuple(),
+        send_args: tuple[Any, ...],
+        send_kwargs: dict[str, Any],
+        retry_codes: set[int] = set(),
+        retry_exceptions: tuple[type[Exception], ...] = tuple(),
     ):
         """
         Send a request and retry it if it fails.
@@ -240,6 +229,11 @@ class PrefectHttpxAsyncClient(httpx.AsyncClient):
         try_count = 0
         response = None
 
+        if TYPE_CHECKING:
+            # older httpx versions type method as str | bytes | Unknown
+            # but in reality it is always a string.
+            assert isinstance(request.method, str)  # type: ignore
+
         is_change_request = request.method.lower() in {"post", "put", "patch", "delete"}
 
         if self.enable_csrf_support and is_change_request:
@@ -297,7 +291,7 @@ class PrefectHttpxAsyncClient(httpx.AsyncClient):
                     if exc_info
                     else (
                         "Received response with retryable status code"
-                        f" {response.status_code}. "
+                        f" {response.status_code if response else 'unknown'}. "
                     )
                 )
                 + f"Another attempt will be made in {retry_seconds}s. "
@@ -314,7 +308,7 @@ class PrefectHttpxAsyncClient(httpx.AsyncClient):
         # We ran out of retries, return the failed response
         return response
 
-    async def send(self, request: Request, *args, **kwargs) -> Response:
+    async def send(self, request: Request, *args: Any, **kwargs: Any) -> Response:
         """
         Send a request with automatic retry behavior for the following status codes:
 
@@ -414,10 +408,10 @@ class PrefectHttpxSyncClient(httpx.Client):
 
     def __init__(
         self,
-        *args,
+        *args: Any,
         enable_csrf_support: bool = False,
         raise_on_all_errors: bool = True,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.enable_csrf_support: bool = enable_csrf_support
         self.csrf_token: Optional[str] = None
@@ -436,10 +430,10 @@ class PrefectHttpxSyncClient(httpx.Client):
         self,
         request: Request,
         send: Callable[[Request], Response],
-        send_args: Tuple,
-        send_kwargs: Dict,
-        retry_codes: Set[int] = set(),
-        retry_exceptions: Tuple[Exception, ...] = tuple(),
+        send_args: tuple[Any, ...],
+        send_kwargs: dict[str, Any],
+        retry_codes: set[int] = set(),
+        retry_exceptions: tuple[type[Exception], ...] = tuple(),
     ):
         """
         Send a request and retry it if it fails.
@@ -454,6 +448,11 @@ class PrefectHttpxSyncClient(httpx.Client):
         try_count = 0
         response = None
 
+        if TYPE_CHECKING:
+            # older httpx versions type method as str | bytes | Unknown
+            # but in reality it is always a string.
+            assert isinstance(request.method, str)  # type: ignore
+
         is_change_request = request.method.lower() in {"post", "put", "patch", "delete"}
 
         if self.enable_csrf_support and is_change_request:
@@ -511,7 +510,7 @@ class PrefectHttpxSyncClient(httpx.Client):
                     if exc_info
                     else (
                         "Received response with retryable status code"
-                        f" {response.status_code}. "
+                        f" {response.status_code if response else 'unknown'}. "
                     )
                 )
                 + f"Another attempt will be made in {retry_seconds}s. "
@@ -528,7 +527,7 @@ class PrefectHttpxSyncClient(httpx.Client):
         # We ran out of retries, return the failed response
         return response
 
-    def send(self, request: Request, *args, **kwargs) -> Response:
+    def send(self, request: Request, *args: Any, **kwargs: Any) -> Response:
         """
         Send a request with automatic retry behavior for the following status codes: