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

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