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

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:
 

prefect/client/cloud.py

@@ -1,11 +1,12 @@
 import re
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, NoReturn, Optional, cast
 from uuid import UUID
 
 import anyio
 import httpx
 import pydantic
 from starlette import status
+from typing_extensions import Self
 
 import prefect.context
 import prefect.settings
@@ -30,7 +31,7 @@ PARSE_API_URL_REGEX = re.compile(r"accounts/(.{36})/workspaces/(.{36})")
 def get_cloud_client(
     host: Optional[str] = None,
     api_key: Optional[str] = None,
-    httpx_settings: Optional[dict] = None,
+    httpx_settings: Optional[dict[str, Any]] = None,
     infer_cloud_url: bool = False,
 ) -> "CloudClient":
     """
@@ -45,6 +46,9 @@ def get_cloud_client(
         configured_url = prefect.settings.PREFECT_API_URL.value()
         host = re.sub(PARSE_API_URL_REGEX, "", configured_url)
 
+    if host is None:
+        raise ValueError("Host was not provided and could not be inferred")
+
     return CloudClient(
         host=host,
         api_key=api_key or PREFECT_API_KEY.value(),
@@ -59,11 +63,14 @@ class CloudUnauthorizedError(PrefectException):
 
 
 class CloudClient:
+    account_id: Optional[str] = None
+    workspace_id: Optional[str] = None
+
     def __init__(
         self,
         host: str,
         api_key: str,
-        httpx_settings: Optional[Dict[str, Any]] = None,
+        httpx_settings: Optional[dict[str, Any]] = None,
     ) -> None:
         httpx_settings = httpx_settings or dict()
         httpx_settings.setdefault("headers", dict())
@@ -76,7 +83,7 @@ class CloudClient:
             **httpx_settings, enable_csrf_support=False
         )
 
-        api_url = prefect.settings.PREFECT_API_URL.value() or ""
+        api_url: str = prefect.settings.PREFECT_API_URL.value() or ""
         if match := (
             re.search(PARSE_API_URL_REGEX, host)
             or re.search(PARSE_API_URL_REGEX, api_url)
@@ -97,7 +104,7 @@ class CloudClient:
 
         return f"{self.account_base_url}/workspaces/{self.workspace_id}"
 
-    async def api_healthcheck(self):
+    async def api_healthcheck(self) -> None:
         """
         Attempts to connect to the Cloud API and raises the encountered exception if not
         successful.
@@ -107,8 +114,8 @@ class CloudClient:
         with anyio.fail_after(10):
             await self.read_workspaces()
 
-    async def read_workspaces(self) -> List[Workspace]:
-        workspaces = pydantic.TypeAdapter(List[Workspace]).validate_python(
+    async def read_workspaces(self) -> list[Workspace]:
+        workspaces = pydantic.TypeAdapter(list[Workspace]).validate_python(
             await self.get("/me/workspaces")
         )
         return workspaces
@@ -121,17 +128,17 @@ class CloudClient:
                 return workspace
         raise ValueError("Current workspace not found")
 
-    async def read_worker_metadata(self) -> Dict[str, Any]:
+    async def read_worker_metadata(self) -> dict[str, Any]:
         response = await self.get(
             f"{self.workspace_base_url}/collections/work_pool_types"
         )
-        return cast(Dict[str, Any], response)
+        return cast(dict[str, Any], response)
 
-    async def read_account_settings(self) -> Dict[str, Any]:
+    async def read_account_settings(self) -> dict[str, Any]:
         response = await self.get(f"{self.account_base_url}/settings")
-        return cast(Dict[str, Any], response)
+        return cast(dict[str, Any], response)
 
-    async def update_account_settings(self, settings: Dict[str, Any]):
+    async def update_account_settings(self, settings: dict[str, Any]) -> None:
         await self.request(
             "PATCH",
             f"{self.account_base_url}/settings",
@@ -142,7 +149,7 @@ class CloudClient:
         response = await self.get(f"{self.account_base_url}/ip_allowlist")
         return IPAllowlist.model_validate(response)
 
-    async def update_account_ip_allowlist(self, updated_allowlist: IPAllowlist):
+    async def update_account_ip_allowlist(self, updated_allowlist: IPAllowlist) -> None:
         await self.request(
             "PUT",
             f"{self.account_base_url}/ip_allowlist",
@@ -172,26 +179,26 @@ class CloudClient:
             json=labels,
         )
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Self:
         await self._client.__aenter__()
         return self
 
-    async def __aexit__(self, *exc_info):
+    async def __aexit__(self, *exc_info: Any) -> None:
         return await self._client.__aexit__(*exc_info)
 
-    def __enter__(self):
+    def __enter__(self) -> NoReturn:
         raise RuntimeError(
             "The `CloudClient` must be entered with an async context. Use 'async "
             "with CloudClient(...)' not 'with CloudClient(...)'"
         )
 
-    def __exit__(self, *_):
+    def __exit__(self, *_: object) -> NoReturn:
         assert False, "This should never be called but must be defined for __enter__"
 
-    async def get(self, route, **kwargs):
+    async def get(self, route: str, **kwargs: Any) -> Any:
         return await self.request("GET", route, **kwargs)
 
-    async def request(self, method, route, **kwargs):
+    async def request(self, method: str, route: str, **kwargs: Any) -> Any:
         try:
             res = await self._client.request(method, route, **kwargs)
             res.raise_for_status()

prefect/client/collections.py

@@ -13,12 +13,12 @@ class CollectionsMetadataClient(Protocol):
     async def __aenter__(self) -> "CollectionsMetadataClient":
         ...
 
-    async def __aexit__(self, *exc_info) -> Any:
+    async def __aexit__(self, *exc_info: Any) -> Any:
         ...
 
 
 def get_collections_metadata_client(
-    httpx_settings: Optional[Dict] = None,
+    httpx_settings: Optional[Dict[str, Any]] = None,
 ) -> "CollectionsMetadataClient":
     """
     Creates a client that can be used to fetch metadata for