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

prefect/utilities/context.py

@@ -1,6 +1,7 @@
+from collections.abc import Generator
 from contextlib import contextmanager
 from contextvars import Context, ContextVar, Token
-from typing import TYPE_CHECKING, Dict, Optional, Tuple, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 from uuid import UUID
 
 if TYPE_CHECKING:
@@ -8,8 +9,8 @@ if TYPE_CHECKING:
 
 
 @contextmanager
-def temporary_context(context: Context):
-    tokens: Dict[ContextVar, Token] = {}
+def temporary_context(context: Context) -> Generator[None, Any, None]:
+    tokens: dict[ContextVar[Any], Token[Any]] = {}
     for key, value in context.items():
         token = key.set(value)
         tokens[key] = token
@@ -38,11 +39,11 @@ def get_flow_run_id() -> Optional[UUID]:
     return None
 
 
-def get_task_and_flow_run_ids() -> Tuple[Optional[UUID], Optional[UUID]]:
+def get_task_and_flow_run_ids() -> tuple[Optional[UUID], Optional[UUID]]:
     """
     Get the task run and flow run ids from the context, if available.
 
     Returns:
-        Tuple[Optional[UUID], Optional[UUID]]: a tuple of the task run id and flow run id
+        tuple[Optional[UUID], Optional[UUID]]: a tuple of the task run id and flow run id
     """
     return get_task_run_id(), get_flow_run_id()

prefect/utilities/dispatch.py

@@ -23,28 +23,39 @@ lookup_type(Base, key) # Foo
 import abc
 import inspect
 import warnings
-from typing import Any, Dict, Optional, Type, TypeVar
+from typing import Any, Literal, Optional, TypeVar, overload
 
-T = TypeVar("T", bound=Type)
+T = TypeVar("T", bound=type[Any])
 
-_TYPE_REGISTRIES: Dict[Type, Dict[str, Type]] = {}
+_TYPE_REGISTRIES: dict[Any, dict[str, Any]] = {}
 
 
-def get_registry_for_type(cls: T) -> Optional[Dict[str, T]]:
+def get_registry_for_type(cls: T) -> Optional[dict[str, T]]:
     """
     Get the first matching registry for a class or any of its base classes.
 
     If not found, `None` is returned.
     """
     return next(
-        filter(
-            lambda registry: registry is not None,
-            (_TYPE_REGISTRIES.get(cls) for cls in cls.mro()),
-        ),
+        (reg for cls in cls.mro() if (reg := _TYPE_REGISTRIES.get(cls)) is not None),
         None,
     )
 
 
+@overload
+def get_dispatch_key(
+    cls_or_instance: Any, allow_missing: Literal[False] = False
+) -> str:
+    ...
+
+
+@overload
+def get_dispatch_key(
+    cls_or_instance: Any, allow_missing: Literal[True] = ...
+) -> Optional[str]:
+    ...
+
+
 def get_dispatch_key(
     cls_or_instance: Any, allow_missing: bool = False
 ) -> Optional[str]:
@@ -89,14 +100,14 @@ def get_dispatch_key(
 
 
 @classmethod
-def _register_subclass_of_base_type(cls, **kwargs):
+def _register_subclass_of_base_type(cls: type[Any], **kwargs: Any) -> None:
     if hasattr(cls, "__init_subclass_original__"):
         cls.__init_subclass_original__(**kwargs)
     elif hasattr(cls, "__pydantic_init_subclass_original__"):
         cls.__pydantic_init_subclass_original__(**kwargs)
 
     # Do not register abstract base classes
-    if abc.ABC in getattr(cls, "__bases__", []):
+    if abc.ABC in cls.__bases__:
         return
 
     register_type(cls)
@@ -123,7 +134,7 @@ def register_base_type(cls: T) -> T:
         cls.__pydantic_init_subclass__ = _register_subclass_of_base_type
     else:
         cls.__init_subclass_original__ = getattr(cls, "__init_subclass__")
-        cls.__init_subclass__ = _register_subclass_of_base_type
+        setattr(cls, "__init_subclass__", _register_subclass_of_base_type)
 
     return cls
 
@@ -190,7 +201,7 @@ def lookup_type(cls: T, dispatch_key: str) -> T:
     Look up a dispatch key in the type registry for the given class.
     """
     # Get the first matching registry for the class or one of its bases
-    registry = get_registry_for_type(cls)
+    registry = get_registry_for_type(cls) or {}
 
     # Look up this type in the registry
     subcls = registry.get(dispatch_key)

prefect/utilities/dockerutils.py

@@ -3,22 +3,12 @@ import os
 import shutil
 import sys
 import warnings
+from collections.abc import Generator, Iterable, Iterator
 from contextlib import contextmanager
 from pathlib import Path, PurePosixPath
 from tempfile import TemporaryDirectory
 from types import TracebackType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Generator,
-    Iterable,
-    List,
-    Optional,
-    TextIO,
-    Tuple,
-    Type,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Optional, TextIO, Union, cast
 from urllib.parse import urlsplit
 
 import pendulum
@@ -29,6 +19,12 @@ import prefect
 from prefect.utilities.importtools import lazy_import
 from prefect.utilities.slugify import slugify
 
+if TYPE_CHECKING:
+    import docker
+    import docker.errors
+    from docker import DockerClient
+    from docker.models.images import Image
+
 CONTAINER_LABELS = {
     "io.prefect.version": prefect.__version__,
 }
@@ -102,10 +98,7 @@ def silence_docker_warnings() -> Generator[None, None, None]:
 # want to have those popping up in various modules and test suites.  Instead,
 # consolidate the imports we need here, and expose them via this module.
 with silence_docker_warnings():
-    if TYPE_CHECKING:
-        import docker
-        from docker import DockerClient
-    else:
+    if not TYPE_CHECKING:
         docker = lazy_import("docker")
 
 
@@ -123,7 +116,8 @@ def docker_client() -> Generator["DockerClient", None, None]:
             "This error is often thrown because Docker is not running. Please ensure Docker is running."
         ) from exc
     finally:
-        client is not None and client.close()
+        if client is not None:
+            client.close()  # type: ignore  # typing stub is not complete
 
 
 class BuildError(Exception):
@@ -207,14 +201,15 @@ class ImageBuilder:
     base_directory: Path
     context: Optional[Path]
     platform: Optional[str]
-    dockerfile_lines: List[str]
+    dockerfile_lines: list[str]
+    temporary_directory: Optional[TemporaryDirectory[str]]
 
     def __init__(
         self,
         base_image: str,
-        base_directory: Path = None,
+        base_directory: Optional[Path] = None,
         platform: Optional[str] = None,
-        context: Path = None,
+        context: Optional[Path] = None,
     ):
         """Create an ImageBuilder
 
@@ -250,7 +245,7 @@ class ImageBuilder:
         return self
 
     def __exit__(
-        self, exc: Type[BaseException], value: BaseException, traceback: TracebackType
+        self, exc: type[BaseException], value: BaseException, traceback: TracebackType
     ) -> None:
         if not self.temporary_directory:
             return
@@ -267,7 +262,9 @@ class ImageBuilder:
         """Add lines to this image's Dockerfile"""
         self.dockerfile_lines.extend(lines)
 
-    def copy(self, source: Union[str, Path], destination: Union[str, PurePosixPath]):
+    def copy(
+        self, source: Union[str, Path], destination: Union[str, PurePosixPath]
+    ) -> None:
         """Copy a file to this image"""
         if not self.context:
             raise Exception("No context available")
@@ -291,7 +288,7 @@ class ImageBuilder:
 
         self.add_line(f"COPY {source} {destination}")
 
-    def write_text(self, text: str, destination: Union[str, PurePosixPath]):
+    def write_text(self, text: str, destination: Union[str, PurePosixPath]) -> None:
         if not self.context:
             raise Exception("No context available")
 
@@ -315,6 +312,7 @@ class ImageBuilder:
         Returns:
             The image ID
         """
+        assert self.context is not None
         dockerfile_path: Path = self.context / "Dockerfile"
 
         with dockerfile_path.open("w") as dockerfile:
@@ -436,9 +434,12 @@ def push_image(
     repository = f"{registry}/{name}"
 
     with docker_client() as client:
-        image: "docker.Image" = client.images.get(image_id)
-        image.tag(repository, tag=tag)
-        events = client.api.push(repository, tag=tag, stream=True, decode=True)
+        image: "Image" = client.images.get(image_id)
+        image.tag(repository, tag=tag)  # type: ignore  # typing stub is not complete
+        events = cast(
+            Iterator[dict[str, Any]],
+            client.api.push(repository, tag=tag, stream=True, decode=True),  # type: ignore  # typing stub is not complete
+        )
         try:
             for event in events:
                 if "status" in event:
@@ -452,12 +453,12 @@ def push_image(
                 elif "error" in event:
                     raise PushError(event["error"])
         finally:
-            client.api.remove_image(f"{repository}:{tag}", noprune=True)
+            client.api.remove_image(f"{repository}:{tag}", noprune=True)  # type: ignore  # typing stub is not complete
 
     return f"{repository}:{tag}"
 
 
-def to_run_command(command: List[str]) -> str:
+def to_run_command(command: list[str]) -> str:
     """
     Convert a process-style list of command arguments to a single Dockerfile RUN
     instruction.
@@ -481,7 +482,7 @@ def to_run_command(command: List[str]) -> str:
     return run_command
 
 
-def parse_image_tag(name: str) -> Tuple[str, Optional[str]]:
+def parse_image_tag(name: str) -> tuple[str, Optional[str]]:
     """
     Parse Docker Image String
 
@@ -519,7 +520,7 @@ def parse_image_tag(name: str) -> Tuple[str, Optional[str]]:
     return image_name, tag
 
 
-def split_repository_path(repository_path: str) -> Tuple[Optional[str], str]:
+def split_repository_path(repository_path: str) -> tuple[Optional[str], str]:
     """
     Splits a Docker repository path into its namespace and repository components.
 
@@ -550,7 +551,7 @@ def split_repository_path(repository_path: str) -> Tuple[Optional[str], str]:
     return namespace, repository
 
 
-def format_outlier_version_name(version: str):
+def format_outlier_version_name(version: str) -> str:
     """
     Formats outlier docker version names to pass `packaging.version.parse` validation
     - Current cases are simple, but creates stub for more complicated formatting if eventually needed.
@@ -580,7 +581,7 @@ def generate_default_dockerfile(context: Optional[Path] = None):
     """
     if not context:
         context = Path.cwd()
-    lines = []
+    lines: list[str] = []
     base_image = get_prefect_image_name()
     lines.append(f"FROM {base_image}")
     dir_name = context.name