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

prefect/utilities/processutils.py

@@ -4,28 +4,35 @@ import signal
 import subprocess
 import sys
 import threading
+from collections.abc import AsyncGenerator, Mapping
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from functools import partial
+from types import FrameType
 from typing import (
     IO,
+    TYPE_CHECKING,
     Any,
+    AnyStr,
     Callable,
-    List,
-    Mapping,
     Optional,
-    Sequence,
     TextIO,
-    Tuple,
     Union,
+    cast,
+    overload,
 )
 
 import anyio
 import anyio.abc
 from anyio.streams.text import TextReceiveStream, TextSendStream
+from typing_extensions import TypeAlias, TypeVar
 
-TextSink = Union[anyio.AsyncFile, TextIO, TextSendStream]
+if TYPE_CHECKING:
+    from _typeshed import StrOrBytesPath
 
+TextSink: TypeAlias = Union[anyio.AsyncFile[AnyStr], TextIO, TextSendStream]
+PrintFn: TypeAlias = Callable[[str], object]
+T = TypeVar("T", infer_variance=True)
 
 if sys.platform == "win32":
     from ctypes import WINFUNCTYPE, c_int, c_uint, windll
@@ -33,7 +40,7 @@ if sys.platform == "win32":
     _windows_process_group_pids = set()
 
     @WINFUNCTYPE(c_int, c_uint)
-    def _win32_ctrl_handler(dwCtrlType):
+    def _win32_ctrl_handler(dwCtrlType: object) -> int:
         """
         A callback function for handling CTRL events cleanly on Windows. When called,
         this function will terminate all running win32 subprocesses the current
@@ -125,16 +132,16 @@ if sys.platform == "win32":
             return self._stderr
 
     async def _open_anyio_process(
-        command: Union[str, bytes, Sequence[Union[str, bytes]]],
+        command: Union[str, bytes, list["StrOrBytesPath"]],
         *,
         stdin: Union[int, IO[Any], None] = None,
         stdout: Union[int, IO[Any], None] = None,
         stderr: Union[int, IO[Any], None] = None,
-        cwd: Union[str, bytes, os.PathLike, None] = None,
-        env: Union[Mapping[str, str], None] = None,
+        cwd: Optional["StrOrBytesPath"] = None,
+        env: Optional[Mapping[str, str]] = None,
         start_new_session: bool = False,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> Process:
         """
         Open a subprocess and return a `Process` object.
 
@@ -179,7 +186,9 @@ if sys.platform == "win32":
 
 
 @asynccontextmanager
-async def open_process(command: List[str], **kwargs):
+async def open_process(
+    command: list[str], **kwargs: Any
+) -> AsyncGenerator[anyio.abc.Process, Any]:
     """
     Like `anyio.open_process` but with:
     - Support for Windows command joining
@@ -189,11 +198,12 @@ async def open_process(command: List[str], **kwargs):
     # Passing a string to open_process is equivalent to shell=True which is
     # generally necessary for Unix-like commands on Windows but otherwise should
     # be avoided
-    if not isinstance(command, list):
-        raise TypeError(
-            "The command passed to open process must be a list. You passed the command"
-            f"'{command}', which is type '{type(command)}'."
-        )
+    if not TYPE_CHECKING:
+        if not isinstance(command, list):
+            raise TypeError(
+                "The command passed to open process must be a list. You passed the command"
+                f"'{command}', which is type '{type(command)}'."
+            )
 
     if sys.platform == "win32":
         command = " ".join(command)
@@ -222,7 +232,7 @@ async def open_process(command: List[str], **kwargs):
     finally:
         try:
             process.terminate()
-            if win32_process_group:
+            if sys.platform == "win32" and win32_process_group:
                 _windows_process_group_pids.remove(process.pid)
 
         except OSError:
@@ -236,13 +246,58 @@ async def open_process(command: List[str], **kwargs):
             await process.aclose()
 
 
+@overload
+async def run_process(
+    command: list[str],
+    *,
+    stream_output: Union[
+        bool, tuple[Optional[TextSink[str]], Optional[TextSink[str]]]
+    ] = ...,
+    task_status: anyio.abc.TaskStatus[T] = ...,
+    task_status_handler: Callable[[anyio.abc.Process], T] = ...,
+    **kwargs: Any,
+) -> anyio.abc.Process:
+    ...
+
+
+@overload
+async def run_process(
+    command: list[str],
+    *,
+    stream_output: Union[
+        bool, tuple[Optional[TextSink[str]], Optional[TextSink[str]]]
+    ] = ...,
+    task_status: Optional[anyio.abc.TaskStatus[int]] = ...,
+    task_status_handler: None = None,
+    **kwargs: Any,
+) -> anyio.abc.Process:
+    ...
+
+
+@overload
+async def run_process(
+    command: list[str],
+    *,
+    stream_output: Union[
+        bool, tuple[Optional[TextSink[str]], Optional[TextSink[str]]]
+    ] = False,
+    task_status: Optional[anyio.abc.TaskStatus[T]] = None,
+    task_status_handler: Optional[Callable[[anyio.abc.Process], T]] = None,
+    **kwargs: Any,
+) -> anyio.abc.Process:
+    ...
+
+
 async def run_process(
-    command: List[str],
-    stream_output: Union[bool, Tuple[Optional[TextSink], Optional[TextSink]]] = False,
-    task_status: Optional[anyio.abc.TaskStatus] = None,
-    task_status_handler: Optional[Callable[[anyio.abc.Process], Any]] = None,
-    **kwargs,
-):
+    command: list[str],
+    *,
+    stream_output: Union[
+        bool, tuple[Optional[TextSink[str]], Optional[TextSink[str]]]
+    ] = False,
+    task_status: Optional[anyio.abc.TaskStatus[T]] = None,
+    task_status_handler: Optional[Callable[[anyio.abc.Process], T]] = None,
+    **kwargs: Any,
+) -> anyio.abc.Process:
     """
     Like `anyio.run_process` but with:
 
@@ -262,12 +317,10 @@ async def run_process(
         **kwargs,
     ) as process:
         if task_status is not None:
-            if not task_status_handler:
-
-                def task_status_handler(process):
-                    return process.pid
-
-            task_status.started(task_status_handler(process))
+            value: T = cast(T, process.pid)
+            if task_status_handler:
+                value = task_status_handler(process)
+            task_status.started(value)
 
         if stream_output:
             await consume_process_output(
@@ -280,31 +333,36 @@ async def run_process(
 
 
 async def consume_process_output(
-    process,
-    stdout_sink: Optional[TextSink] = None,
-    stderr_sink: Optional[TextSink] = None,
-):
+    process: anyio.abc.Process,
+    stdout_sink: Optional[TextSink[str]] = None,
+    stderr_sink: Optional[TextSink[str]] = None,
+) -> None:
     async with anyio.create_task_group() as tg:
-        tg.start_soon(
-            stream_text,
-            TextReceiveStream(process.stdout),
-            stdout_sink,
-        )
-        tg.start_soon(
-            stream_text,
-            TextReceiveStream(process.stderr),
-            stderr_sink,
-        )
+        if process.stdout is not None:
+            tg.start_soon(
+                stream_text,
+                TextReceiveStream(process.stdout),
+                stdout_sink,
+            )
+        if process.stderr is not None:
+            tg.start_soon(
+                stream_text,
+                TextReceiveStream(process.stderr),
+                stderr_sink,
+            )
 
 
-async def stream_text(source: TextReceiveStream, *sinks: TextSink):
+async def stream_text(
+    source: TextReceiveStream, *sinks: Optional[TextSink[str]]
+) -> None:
     wrapped_sinks = [
         (
-            anyio.wrap_file(sink)
+            anyio.wrap_file(cast(IO[str], sink))
             if hasattr(sink, "write") and hasattr(sink, "flush")
             else sink
         )
         for sink in sinks
+        if sink is not None
     ]
     async for item in source:
         for sink in wrapped_sinks:
@@ -313,30 +371,32 @@ async def stream_text(source: TextReceiveStream, *sinks: TextSink):
             elif isinstance(sink, anyio.AsyncFile):
                 await sink.write(item)
                 await sink.flush()
-            elif sink is None:
-                pass  # Consume the item but perform no action
-            else:
-                raise TypeError(f"Unsupported sink type {type(sink).__name__}")
 
 
-def _register_signal(signum: int, handler: Callable):
+def _register_signal(
+    signum: int,
+    handler: Optional[
+        Union[Callable[[int, Optional[FrameType]], Any], int, signal.Handlers]
+    ],
+) -> None:
     if threading.current_thread() is threading.main_thread():
         signal.signal(signum, handler)
 
 
 def forward_signal_handler(
-    pid: int, signum: int, *signums: int, process_name: str, print_fn: Callable
-):
+    pid: int, signum: int, *signums: int, process_name: str, print_fn: PrintFn
+) -> None:
     """Forward subsequent signum events (e.g. interrupts) to respective signums."""
     current_signal, future_signals = signums[0], signums[1:]
 
     # avoid RecursionError when setting up a direct signal forward to the same signal for the main pid
+    original_handler = None
     avoid_infinite_recursion = signum == current_signal and pid == os.getpid()
     if avoid_infinite_recursion:
         # store the vanilla handler so it can be temporarily restored below
         original_handler = signal.getsignal(current_signal)
 
-    def handler(*args):
+    def handler(*arg: Any) -> None:
         print_fn(
             f"Received {getattr(signum, 'name', signum)}. "
             f"Sending {getattr(current_signal, 'name', current_signal)} to"
@@ -358,7 +418,9 @@ def forward_signal_handler(
     _register_signal(signum, handler)
 
 
-def setup_signal_handlers_server(pid: int, process_name: str, print_fn: Callable):
+def setup_signal_handlers_server(
+    pid: int, process_name: str, print_fn: PrintFn
+) -> None:
     """Handle interrupts of the server gracefully."""
     setup_handler = partial(
         forward_signal_handler, pid, process_name=process_name, print_fn=print_fn
@@ -375,7 +437,7 @@ def setup_signal_handlers_server(pid: int, process_name: str, print_fn: Callable
         setup_handler(signal.SIGTERM, signal.SIGTERM, signal.SIGKILL)
 
 
-def setup_signal_handlers_agent(pid: int, process_name: str, print_fn: Callable):
+def setup_signal_handlers_agent(pid: int, process_name: str, print_fn: PrintFn) -> None:
     """Handle interrupts of the agent gracefully."""
     setup_handler = partial(
         forward_signal_handler, pid, process_name=process_name, print_fn=print_fn
@@ -393,7 +455,9 @@ def setup_signal_handlers_agent(pid: int, process_name: str, print_fn: Callable)
         setup_handler(signal.SIGTERM, signal.SIGINT, signal.SIGKILL)
 
 
-def setup_signal_handlers_worker(pid: int, process_name: str, print_fn: Callable):
+def setup_signal_handlers_worker(
+    pid: int, process_name: str, print_fn: PrintFn
+) -> None:
     """Handle interrupts of workers gracefully."""
     setup_handler = partial(
         forward_signal_handler, pid, process_name=process_name, print_fn=print_fn

prefect/utilities/pydantic.py

@@ -1,18 +1,18 @@
-from functools import partial
 from typing import (
     Any,
     Callable,
-    Dict,
     Generic,
     Optional,
-    Type,
     TypeVar,
+    Union,
     cast,
     get_origin,
     overload,
 )
 
-from jsonpatch import JsonPatch as JsonPatchBase
+from jsonpatch import (  # type: ignore  # no typing stubs available, see https://github.com/stefankoegl/python-json-patch/issues/158
+    JsonPatch as JsonPatchBase,
+)
 from pydantic import (
     BaseModel,
     GetJsonSchemaHandler,
@@ -33,7 +33,7 @@ M = TypeVar("M", bound=BaseModel)
 T = TypeVar("T", bound=Any)
 
 
-def _reduce_model(model: BaseModel):
+def _reduce_model(self: BaseModel) -> tuple[Any, ...]:
     """
     Helper for serializing a cythonized model with cloudpickle.
 
@@ -43,31 +43,33 @@ def _reduce_model(model: BaseModel):
     return (
         _unreduce_model,
         (
-            to_qualified_name(type(model)),
-            model.model_dump_json(**getattr(model, "__reduce_kwargs__", {})),
+            to_qualified_name(type(self)),
+            self.model_dump_json(**getattr(self, "__reduce_kwargs__", {})),
         ),
     )
 
 
-def _unreduce_model(model_name, json):
+def _unreduce_model(model_name: str, json: str) -> Any:
     """Helper for restoring model after serialization"""
     model = from_qualified_name(model_name)
     return model.model_validate_json(json)
 
 
 @overload
-def add_cloudpickle_reduction(__model_cls: Type[M]) -> Type[M]:
+def add_cloudpickle_reduction(__model_cls: type[M]) -> type[M]:
     ...
 
 
 @overload
 def add_cloudpickle_reduction(
-    **kwargs: Any,
-) -> Callable[[Type[M]], Type[M]]:
+    __model_cls: None = None, **kwargs: Any
+) -> Callable[[type[M]], type[M]]:
     ...
 
 
-def add_cloudpickle_reduction(__model_cls: Optional[Type[M]] = None, **kwargs: Any):
+def add_cloudpickle_reduction(
+    __model_cls: Optional[type[M]] = None, **kwargs: Any
+) -> Union[type[M], Callable[[type[M]], type[M]]]:
     """
     Adds a `__reducer__` to the given class that ensures it is cloudpickle compatible.
 
@@ -85,25 +87,22 @@ def add_cloudpickle_reduction(__model_cls: Optional[Type[M]] = None, **kwargs: A
     """
     if __model_cls:
         __model_cls.__reduce__ = _reduce_model
-        __model_cls.__reduce_kwargs__ = kwargs
+        setattr(__model_cls, "__reduce_kwargs__", kwargs)
         return __model_cls
-    else:
-        return cast(
-            Callable[[Type[M]], Type[M]],
-            partial(
-                add_cloudpickle_reduction,
-                **kwargs,
-            ),
-        )
+
+    def reducer_with_kwargs(__model_cls: type[M]) -> type[M]:
+        return add_cloudpickle_reduction(__model_cls, **kwargs)
+
+    return reducer_with_kwargs
 
 
-def get_class_fields_only(model: Type[BaseModel]) -> set:
+def get_class_fields_only(model: type[BaseModel]) -> set[str]:
     """
     Gets all the field names defined on the model class but not any parent classes.
     Any fields that are on the parent but redefined on the subclass are included.
     """
     subclass_class_fields = set(model.__annotations__.keys())
-    parent_class_fields = set()
+    parent_class_fields: set[str] = set()
 
     for base in model.__class__.__bases__:
         if issubclass(base, BaseModel):
@@ -114,7 +113,7 @@ def get_class_fields_only(model: Type[BaseModel]) -> set:
     )
 
 
-def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
+def add_type_dispatch(model_cls: type[M]) -> type[M]:
     """
     Extend a Pydantic model to add a 'type' field that is used as a discriminator field
     to dynamically determine the subtype that when deserializing models.
@@ -149,7 +148,7 @@ def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
 
     elif not defines_dispatch_key and defines_type_field:
         field_type_annotation = model_cls.model_fields["type"].annotation
-        if field_type_annotation != str:
+        if field_type_annotation != str and field_type_annotation is not None:
             raise TypeError(
                 f"Model class {model_cls.__name__!r} defines a 'type' field with "
                 f"type {field_type_annotation.__name__!r} but it must be 'str'."
@@ -157,10 +156,10 @@ def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
 
         # Set the dispatch key to retrieve the value from the type field
         @classmethod
-        def dispatch_key_from_type_field(cls):
+        def dispatch_key_from_type_field(cls: type[M]) -> str:
             return cls.model_fields["type"].default
 
-        model_cls.__dispatch_key__ = dispatch_key_from_type_field
+        setattr(model_cls, "__dispatch_key__", dispatch_key_from_type_field)
 
     else:
         raise ValueError(
@@ -171,7 +170,7 @@ def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
     cls_init = model_cls.__init__
     cls_new = model_cls.__new__
 
-    def __init__(__pydantic_self__, **data: Any) -> None:
+    def __init__(__pydantic_self__: M, **data: Any) -> None:
         type_string = (
             get_dispatch_key(__pydantic_self__)
             if type(__pydantic_self__) != model_cls
@@ -180,12 +179,16 @@ def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
         data.setdefault("type", type_string)
         cls_init(__pydantic_self__, **data)
 
-    def __new__(cls: Type[M], **kwargs: Any) -> M:
+    def __new__(cls: type[M], **kwargs: Any) -> M:
         if "type" in kwargs:
             try:
                 subcls = lookup_type(cls, dispatch_key=kwargs["type"])
             except KeyError as exc:
-                raise ValidationError(errors=[exc], model=cls)
+                raise ValidationError.from_exception_data(
+                    title=cls.__name__,
+                    line_errors=[{"type": str(exc), "input": kwargs["type"]}],
+                    input_type="python",
+                )
             return cls_new(subcls)
         else:
             return cls_new(cls)
@@ -221,7 +224,7 @@ class PartialModel(Generic[M]):
         >>> model = partial_model.finalize(z=3.0)
     """
 
-    def __init__(self, __model_cls: Type[M], **kwargs: Any) -> None:
+    def __init__(self, __model_cls: type[M], **kwargs: Any) -> None:
         self.fields = kwargs
         # Set fields first to avoid issues if `fields` is also set on the `model_cls`
         # in our custom `setattr` implementation.
@@ -236,11 +239,11 @@ class PartialModel(Generic[M]):
             self.raise_if_not_in_model(name)
         return self.model_cls(**self.fields, **kwargs)
 
-    def raise_if_already_set(self, name):
+    def raise_if_already_set(self, name: str) -> None:
         if name in self.fields:
             raise ValueError(f"Field {name!r} has already been set.")
 
-    def raise_if_not_in_model(self, name):
+    def raise_if_not_in_model(self, name: str) -> None:
         if name not in self.model_cls.model_fields:
             raise ValueError(f"Field {name!r} is not present in the model.")
 
@@ -290,7 +293,7 @@ class JsonPatch(JsonPatchBase):
 
 
 def custom_pydantic_encoder(
-    type_encoders: Optional[Dict[Any, Callable[[Type[Any]], Any]]], obj: Any
+    type_encoders: dict[Any, Callable[[type[Any]], Any]], obj: Any
 ) -> Any:
     # Check the class type and its superclasses for a matching encoder
     for base in obj.__class__.__mro__[:-1]:
@@ -359,8 +362,10 @@ def parse_obj_as(
     """
     adapter = TypeAdapter(type_)
 
-    if get_origin(type_) is list and isinstance(data, dict):
-        data = next(iter(data.values()))
+    origin: Optional[Any] = get_origin(type_)
+    if origin is list and isinstance(data, dict):
+        values_dict: dict[Any, Any] = data
+        data = next(iter(values_dict.values()))
 
     parser: Callable[[Any], T] = getattr(adapter, f"validate_{mode}")
 

prefect/utilities/render_swagger.py

@@ -8,10 +8,13 @@ import re
 import string
 import urllib.parse
 from pathlib import Path
+from typing import Any, Optional, cast
 from xml.sax.saxutils import escape
 
 import mkdocs.plugins
-from mkdocs.structure.files import File
+from mkdocs.config.defaults import MkDocsConfig
+from mkdocs.structure.files import File, Files
+from mkdocs.structure.pages import Page
 
 USAGE_MSG = (
     "Usage: '!!swagger <filename>!!' or '!!swagger-http <url>!!'. "
@@ -50,7 +53,7 @@ TOKEN = re.compile(r"!!swagger(?: (?P<path>[^\\/\s><&:]+))?!!")
 TOKEN_HTTP = re.compile(r"!!swagger-http(?: (?P<path>https?://[^\s]+))?!!")
 
 
-def swagger_lib(config) -> dict:
+def swagger_lib(config: MkDocsConfig) -> dict[str, Any]:
     """
     Provides the actual swagger library used
     """
@@ -59,11 +62,14 @@ def swagger_lib(config) -> dict:
         "js": "https://unpkg.com/swagger-ui-dist@3/swagger-ui-bundle.js",
     }
 
-    extra_javascript = config.get("extra_javascript", [])
-    extra_css = config.get("extra_css", [])
+    extra_javascript = config.extra_javascript
+    extra_css = cast(list[str], config.extra_css)
     for lib in extra_javascript:
-        if os.path.basename(urllib.parse.urlparse(lib).path) == "swagger-ui-bundle.js":
-            lib_swagger["js"] = lib
+        if (
+            os.path.basename(urllib.parse.urlparse(str(lib)).path)
+            == "swagger-ui-bundle.js"
+        ):
+            lib_swagger["js"] = str(lib)
             break
 
     for css in extra_css:
@@ -73,8 +79,10 @@ def swagger_lib(config) -> dict:
     return lib_swagger
 
 
-class SwaggerPlugin(mkdocs.plugins.BasePlugin):
-    def on_page_markdown(self, markdown, page, config, files):
+class SwaggerPlugin(mkdocs.plugins.BasePlugin[Any]):
+    def on_page_markdown(
+        self, markdown: str, /, *, page: Page, config: MkDocsConfig, files: Files
+    ) -> Optional[str]:
         is_http = False
         match = TOKEN.search(markdown)
 
@@ -88,7 +96,7 @@ class SwaggerPlugin(mkdocs.plugins.BasePlugin):
         pre_token = markdown[: match.start()]
         post_token = markdown[match.end() :]
 
-        def _error(message):
+        def _error(message: str) -> str:
             return (
                 pre_token
                 + escape(ERROR_TEMPLATE.substitute(error=message))
@@ -103,8 +111,10 @@ class SwaggerPlugin(mkdocs.plugins.BasePlugin):
         if is_http:
             url = path
         else:
+            base = page.file.abs_src_path
+            assert base is not None
             try:
-                api_file = Path(page.file.abs_src_path).with_name(path)
+                api_file = Path(base).with_name(path)
             except ValueError as exc:
                 return _error(f"Invalid path. {exc.args[0]}")
 
@@ -114,7 +124,7 @@ class SwaggerPlugin(mkdocs.plugins.BasePlugin):
             src_dir = api_file.parent
             dest_dir = Path(page.file.abs_dest_path).parent
 
-            new_file = File(api_file.name, src_dir, dest_dir, False)
+            new_file = File(api_file.name, str(src_dir), str(dest_dir), False)
             files.append(new_file)
             url = Path(new_file.abs_dest_path).name
 
@@ -129,4 +139,4 @@ class SwaggerPlugin(mkdocs.plugins.BasePlugin):
         )
 
         # If multiple swaggers exist.
-        return self.on_page_markdown(markdown, page, config, files)
+        return self.on_page_markdown(markdown, page=page, config=config, files=files)

prefect/utilities/schema_tools/__init__.py

@@ -2,8 +2,8 @@ from .hydration import HydrationContext, HydrationError, hydrate
 from .validation import (
     CircularSchemaRefError,
     ValidationError,
-    validate,
     is_valid_schema,
+    validate,
 )
 
 __all__ = [
@@ -12,5 +12,6 @@ __all__ = [
     "HydrationError",
     "ValidationError",
     "hydrate",
+    "is_valid_schema",
     "validate",
 ]