prefect-client 3.1.11__py3-none-any.whl → 3.1.13__py3-none-any.whl

prefect/input/run_input.py

@@ -60,11 +60,15 @@ async def receiver_flow():
 ```
 """
 
+from __future__ import annotations
+
+import inspect
 from inspect import isclass
 from typing import (
     TYPE_CHECKING,
     Any,
     ClassVar,
+    Coroutine,
     Dict,
     Generic,
     Literal,
@@ -81,6 +85,7 @@ from uuid import UUID, uuid4
 import anyio
 import pydantic
 from pydantic import ConfigDict
+from typing_extensions import Self
 
 from prefect.input.actions import (
     create_flow_run_input,
@@ -144,7 +149,7 @@ class RunInputMetadata(pydantic.BaseModel):
     receiver: UUID
 
 
-class RunInput(pydantic.BaseModel):
+class BaseRunInput(pydantic.BaseModel):
     model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
 
     _description: Optional[str] = pydantic.PrivateAttr(default=None)
@@ -172,23 +177,29 @@ class RunInput(pydantic.BaseModel):
         if is_v2_model(cls):
             schema = create_v2_schema(cls.__name__, model_base=cls)
         else:
-            schema = cls.schema(by_alias=True)
+            schema = cls.model_json_schema(by_alias=True)
 
-        await create_flow_run_input(
+        coro = create_flow_run_input(
             key=keyset["schema"], value=schema, flow_run_id=flow_run_id
         )
+        if TYPE_CHECKING:
+            assert inspect.iscoroutine(coro)
+        await coro
 
         description = cls._description if isinstance(cls._description, str) else None
         if description:
-            await create_flow_run_input(
+            coro = create_flow_run_input(
                 key=keyset["description"],
                 value=description,
                 flow_run_id=flow_run_id,
             )
+            if TYPE_CHECKING:
+                assert inspect.iscoroutine(coro)
+            await coro
 
     @classmethod
     @sync_compatible
-    async def load(cls, keyset: Keyset, flow_run_id: Optional[UUID] = None):
+    async def load(cls, keyset: Keyset, flow_run_id: Optional[UUID] = None) -> Self:
         """
         Load the run input response from the given key.
 
@@ -208,7 +219,7 @@ class RunInput(pydantic.BaseModel):
         return instance
 
     @classmethod
-    def load_from_flow_run_input(cls, flow_run_input: "FlowRunInput"):
+    def load_from_flow_run_input(cls, flow_run_input: "FlowRunInput") -> Self:
         """
         Load the run input from a FlowRunInput object.
 
@@ -284,6 +295,8 @@ class RunInput(pydantic.BaseModel):
             key_prefix=key_prefix,
         )
 
+
+class RunInput(BaseRunInput):
     @classmethod
     def receive(
         cls,
@@ -293,7 +306,7 @@ class RunInput(pydantic.BaseModel):
         exclude_keys: Optional[Set[str]] = None,
         key_prefix: Optional[str] = None,
         flow_run_id: Optional[UUID] = None,
-    ):
+    ) -> GetInputHandler[Self]:
         if key_prefix is None:
             key_prefix = f"{cls.__name__.lower()}-auto"
 
@@ -322,12 +335,12 @@ class RunInput(pydantic.BaseModel):
         return type(f"{model_cls.__name__}RunInput", (RunInput, model_cls), {})  # type: ignore
 
 
-class AutomaticRunInput(RunInput, Generic[T]):
+class AutomaticRunInput(BaseRunInput, Generic[T]):
     value: T
 
     @classmethod
     @sync_compatible
-    async def load(cls, keyset: Keyset, flow_run_id: Optional[UUID] = None) -> T:
+    async def load(cls, keyset: Keyset, flow_run_id: Optional[UUID] = None) -> Self:
         """
         Load the run input response from the given key.
 
@@ -335,7 +348,10 @@ class AutomaticRunInput(RunInput, Generic[T]):
             - keyset (Keyset): the keyset to load the input for
             - flow_run_id (UUID, optional): the flow run ID to load the input for
         """
-        instance = await super().load(keyset, flow_run_id=flow_run_id)
+        instance_coro = super().load(keyset, flow_run_id=flow_run_id)
+        if TYPE_CHECKING:
+            assert inspect.iscoroutine(instance_coro)
+        instance = await instance_coro
         return instance.value
 
     @classmethod
@@ -370,17 +386,34 @@ class AutomaticRunInput(RunInput, Generic[T]):
 
         # Creating a new Pydantic model class dynamically with the name based
         # on the type prefix.
-        new_cls: Type["AutomaticRunInput"] = pydantic.create_model(
+        new_cls: Type["AutomaticRunInput[T]"] = pydantic.create_model(
             class_name, **fields, __base__=AutomaticRunInput
         )
         return new_cls
 
     @classmethod
-    def receive(cls, *args, **kwargs):
-        if kwargs.get("key_prefix") is None:
-            kwargs["key_prefix"] = f"{cls.__name__.lower()}-auto"
+    def receive(
+        cls,
+        timeout: Optional[float] = 3600,
+        poll_interval: float = 10,
+        raise_timeout_error: bool = False,
+        exclude_keys: Optional[Set[str]] = None,
+        key_prefix: Optional[str] = None,
+        flow_run_id: Optional[UUID] = None,
+        with_metadata: bool = False,
+    ) -> GetAutomaticInputHandler[T]:
+        key_prefix = key_prefix or f"{cls.__name__.lower()}-auto"
 
-        return GetAutomaticInputHandler(run_input_cls=cls, *args, **kwargs)
+        return GetAutomaticInputHandler(
+            run_input_cls=cls,
+            key_prefix=key_prefix,
+            timeout=timeout,
+            poll_interval=poll_interval,
+            raise_timeout_error=raise_timeout_error,
+            exclude_keys=exclude_keys,
+            flow_run_id=flow_run_id,
+            with_metadata=with_metadata,
+        )
 
 
 def run_input_subclass_from_type(
@@ -409,24 +442,24 @@ class GetInputHandler(Generic[R]):
         self,
         run_input_cls: Type[R],
         key_prefix: str,
-        timeout: Optional[float] = 3600,
+        timeout: float | None = 3600,
         poll_interval: float = 10,
         raise_timeout_error: bool = False,
         exclude_keys: Optional[Set[str]] = None,
         flow_run_id: Optional[UUID] = None,
     ):
-        self.run_input_cls = run_input_cls
-        self.key_prefix = key_prefix
-        self.timeout = timeout
-        self.poll_interval = poll_interval
-        self.exclude_keys = set()
-        self.raise_timeout_error = raise_timeout_error
-        self.flow_run_id = ensure_flow_run_id(flow_run_id)
+        self.run_input_cls: Type[R] = run_input_cls
+        self.key_prefix: str = key_prefix
+        self.timeout: float | None = timeout
+        self.poll_interval: float = poll_interval
+        self.exclude_keys: set[str] = set()
+        self.raise_timeout_error: bool = raise_timeout_error
+        self.flow_run_id: UUID = ensure_flow_run_id(flow_run_id)
 
         if exclude_keys is not None:
             self.exclude_keys.update(exclude_keys)
 
-    def __iter__(self):
+    def __iter__(self) -> Self:
         return self
 
     def __next__(self) -> R:
@@ -437,24 +470,31 @@ class GetInputHandler(Generic[R]):
                 raise
             raise StopIteration
 
-    def __aiter__(self):
+    def __aiter__(self) -> Self:
         return self
 
     async def __anext__(self) -> R:
         try:
-            return await self.next()
+            coro = self.next()
+            if TYPE_CHECKING:
+                assert inspect.iscoroutine(coro)
+            return await coro
         except TimeoutError:
             if self.raise_timeout_error:
                 raise
             raise StopAsyncIteration
 
-    async def filter_for_inputs(self):
-        flow_run_inputs = await filter_flow_run_input(
+    async def filter_for_inputs(self) -> list["FlowRunInput"]:
+        flow_run_inputs_coro = filter_flow_run_input(
             key_prefix=self.key_prefix,
             limit=1,
             exclude_keys=self.exclude_keys,
             flow_run_id=self.flow_run_id,
         )
+        if TYPE_CHECKING:
+            assert inspect.iscoroutine(flow_run_inputs_coro)
+
+        flow_run_inputs = await flow_run_inputs_coro
 
         if flow_run_inputs:
             self.exclude_keys.add(*[i.key for i in flow_run_inputs])
@@ -478,22 +518,91 @@ class GetInputHandler(Generic[R]):
                     return self.to_instance(flow_run_inputs[0])
 
 
-class GetAutomaticInputHandler(GetInputHandler, Generic[T]):
-    def __init__(self, *args, **kwargs):
-        self.with_metadata = kwargs.pop("with_metadata", False)
-        super().__init__(*args, **kwargs)
+class GetAutomaticInputHandler(Generic[T]):
+    def __init__(
+        self,
+        run_input_cls: Type[AutomaticRunInput[T]],
+        key_prefix: str,
+        timeout: float | None = 3600,
+        poll_interval: float = 10,
+        raise_timeout_error: bool = False,
+        exclude_keys: Optional[Set[str]] = None,
+        flow_run_id: Optional[UUID] = None,
+        with_metadata: bool = False,
+    ):
+        self.run_input_cls: Type[AutomaticRunInput[T]] = run_input_cls
+        self.key_prefix: str = key_prefix
+        self.timeout: float | None = timeout
+        self.poll_interval: float = poll_interval
+        self.exclude_keys: set[str] = set()
+        self.raise_timeout_error: bool = raise_timeout_error
+        self.flow_run_id: UUID = ensure_flow_run_id(flow_run_id)
+        self.with_metadata = with_metadata
 
-    def __next__(self) -> T:
-        return cast(T, super().__next__())
+        if exclude_keys is not None:
+            self.exclude_keys.update(exclude_keys)
 
-    async def __anext__(self) -> T:
-        return cast(T, await super().__anext__())
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> T | AutomaticRunInput[T]:
+        try:
+            not_coro = self.next()
+            if TYPE_CHECKING:
+                assert not isinstance(not_coro, Coroutine)
+            return not_coro
+        except TimeoutError:
+            if self.raise_timeout_error:
+                raise
+            raise StopIteration
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> Union[T, AutomaticRunInput[T]]:
+        try:
+            coro = self.next()
+            if TYPE_CHECKING:
+                assert inspect.iscoroutine(coro)
+            return cast(Union[T, AutomaticRunInput[T]], await coro)
+        except TimeoutError:
+            if self.raise_timeout_error:
+                raise
+            raise StopAsyncIteration
+
+    async def filter_for_inputs(self) -> list["FlowRunInput"]:
+        flow_run_inputs_coro = filter_flow_run_input(
+            key_prefix=self.key_prefix,
+            limit=1,
+            exclude_keys=self.exclude_keys,
+            flow_run_id=self.flow_run_id,
+        )
+        if TYPE_CHECKING:
+            assert inspect.iscoroutine(flow_run_inputs_coro)
+
+        flow_run_inputs = await flow_run_inputs_coro
+
+        if flow_run_inputs:
+            self.exclude_keys.add(*[i.key for i in flow_run_inputs])
+
+        return flow_run_inputs
 
     @sync_compatible
-    async def next(self) -> T:
-        return cast(T, await super().next())
+    async def next(self) -> Union[T, AutomaticRunInput[T]]:
+        flow_run_inputs = await self.filter_for_inputs()
+        if flow_run_inputs:
+            return self.to_instance(flow_run_inputs[0])
 
-    def to_instance(self, flow_run_input: "FlowRunInput") -> T:
+        with anyio.fail_after(self.timeout):
+            while True:
+                await anyio.sleep(self.poll_interval)
+                flow_run_inputs = await self.filter_for_inputs()
+                if flow_run_inputs:
+                    return self.to_instance(flow_run_inputs[0])
+
+    def to_instance(
+        self, flow_run_input: "FlowRunInput"
+    ) -> Union[T, AutomaticRunInput[T]]:
         run_input = self.run_input_cls.load_from_flow_run_input(flow_run_input)
 
         if self.with_metadata:
@@ -503,14 +612,15 @@ class GetAutomaticInputHandler(GetInputHandler, Generic[T]):
 
 async def _send_input(
     flow_run_id: UUID,
-    run_input: Any,
+    run_input: RunInput | pydantic.BaseModel,
     sender: Optional[str] = None,
     key_prefix: Optional[str] = None,
 ):
+    _run_input: Union[RunInput, AutomaticRunInput[Any]]
     if isinstance(run_input, RunInput):
-        _run_input: RunInput = run_input
+        _run_input = run_input
     else:
-        input_cls: Type[AutomaticRunInput] = run_input_subclass_from_type(
+        input_cls: Type[AutomaticRunInput[Any]] = run_input_subclass_from_type(
             type(run_input)
         )
         _run_input = input_cls(value=run_input)
@@ -520,9 +630,13 @@ async def _send_input(
 
     key = f"{key_prefix}-{uuid4()}"
 
-    await create_flow_run_input_from_model(
+    coro = create_flow_run_input_from_model(
         key=key, flow_run_id=flow_run_id, model_instance=_run_input, sender=sender
     )
+    if TYPE_CHECKING:
+        assert inspect.iscoroutine(coro)
+
+    await coro
 
 
 @sync_compatible

prefect/logging/configuration.py

@@ -6,7 +6,7 @@ import string
 import warnings
 from functools import partial
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Dict, Optional
 
 import yaml
 
@@ -24,10 +24,10 @@ DEFAULT_LOGGING_SETTINGS_PATH = Path(__file__).parent / "logging.yml"
 PROCESS_LOGGING_CONFIG: Optional[Dict[str, Any]] = None
 
 # Regex call to replace non-alphanumeric characters to '_' to create a valid env var
-to_envvar = partial(re.sub, re.compile(r"[^0-9a-zA-Z]+"), "_")
+to_envvar: Callable[[str], str] = partial(re.sub, re.compile(r"[^0-9a-zA-Z]+"), "_")
 
 
-def load_logging_config(path: Path) -> dict:
+def load_logging_config(path: Path) -> dict[str, Any]:
     """
     Loads logging configuration from a path allowing override from the environment
     """

prefect/logging/filters.py

@@ -5,7 +5,7 @@ from prefect.utilities.collections import visit_collection
 from prefect.utilities.names import obfuscate
 
 
-def redact_substr(obj: Any, substr: str):
+def redact_substr(obj: Any, substr: str) -> Any:
     """
     Redact a string from a potentially nested object.
 
@@ -17,7 +17,7 @@ def redact_substr(obj: Any, substr: str):
         Any: The object with the API key redacted.
     """
 
-    def redact_item(item):
+    def redact_item(item: Any) -> Any:
         if isinstance(item, str):
             return item.replace(substr, obfuscate(substr))
         return item

prefect/logging/formatters.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import logging.handlers
 import sys
 import traceback
 from types import TracebackType
-from typing import Optional, Tuple, Type, Union
+from typing import Any, Literal, Optional, Tuple, Type, Union
 
 import orjson
 
@@ -14,7 +16,7 @@ ExceptionInfoType = Union[
 ]
 
 
-def format_exception_info(exc_info: ExceptionInfoType) -> dict:
+def format_exception_info(exc_info: ExceptionInfoType) -> dict[str, Any]:
     # if sys.exc_info() returned a (None, None, None) tuple,
     # then there's nothing to format
     if exc_info[0] is None:
@@ -40,13 +42,15 @@ class JsonFormatter(logging.Formatter):
     newlines.
     """
 
-    def __init__(self, fmt, dmft, style) -> None:  # noqa
+    def __init__(
+        self, fmt: Literal["pretty", "default"], dmft: str, style: str
+    ) -> None:  # noqa
         super().__init__()
 
         if fmt not in ["pretty", "default"]:
             raise ValueError("Format must be either 'pretty' or 'default'.")
 
-        self.serializer = JSONSerializer(
+        self.serializer: JSONSerializer = JSONSerializer(
             jsonlib="orjson",
             dumps_kwargs={"option": orjson.OPT_INDENT_2} if fmt == "pretty" else {},
         )
@@ -72,13 +76,13 @@ class JsonFormatter(logging.Formatter):
 class PrefectFormatter(logging.Formatter):
     def __init__(
         self,
-        format=None,
-        datefmt=None,
-        style="%",
-        validate=True,
+        format: str | None = None,
+        datefmt: str | None = None,
+        style: str = "%",
+        validate: bool = True,
         *,
-        defaults=None,
-        task_run_fmt: Optional[str] = None,
+        defaults: dict[str, Any] | None = None,
+        task_run_fmt: str | None = None,
         flow_run_fmt: Optional[str] = None,
     ) -> None:
         """
@@ -118,7 +122,7 @@ class PrefectFormatter(logging.Formatter):
             self._flow_run_style.validate()
             self._task_run_style.validate()
 
-    def formatMessage(self, record: logging.LogRecord):
+    def formatMessage(self, record: logging.LogRecord) -> str:
         if record.name == "prefect.flow_runs":
             style = self._flow_run_style
         elif record.name == "prefect.task_runs":

prefect/logging/handlers.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import logging
 import sys
@@ -6,7 +8,7 @@ import traceback
 import uuid
 import warnings
 from contextlib import asynccontextmanager
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, TextIO, Type
 
 import pendulum
 from rich.console import Console
@@ -34,6 +36,14 @@ from prefect.settings import (
     PREFECT_LOGGING_TO_API_WHEN_MISSING_FLOW,
 )
 
+if sys.version_info >= (3, 12):
+    StreamHandler = logging.StreamHandler[TextIO]
+else:
+    if TYPE_CHECKING:
+        StreamHandler = logging.StreamHandler[TextIO]
+    else:
+        StreamHandler = logging.StreamHandler
+
 
 class APILogWorker(BatchedQueueService[Dict[str, Any]]):
     @property
@@ -90,7 +100,7 @@ class APILogHandler(logging.Handler):
     """
 
     @classmethod
-    def flush(cls):
+    def flush(cls) -> None:
         """
         Tell the `APILogWorker` to send any currently enqueued logs and block until
         completion.
@@ -107,8 +117,8 @@ class APILogHandler(logging.Handler):
                 )
 
             # Not ideal, but this method is called by the stdlib and cannot return a
-            # coroutine so we just schedule the drain in the global loop thread and continue
-            from_sync.call_soon_in_loop_thread(create_call(APILogWorker.drain_all))
+            # coroutine so we just schedule the drain in a new thread and continue
+            from_sync.call_soon_in_new_thread(create_call(APILogWorker.drain_all))
             return None
         else:
             # We set a timeout of 5s because we don't want to block forever if the worker
@@ -118,7 +128,7 @@ class APILogHandler(logging.Handler):
             return APILogWorker.drain_all(timeout=5)
 
     @classmethod
-    async def aflush(cls):
+    async def aflush(cls) -> bool:
         """
         Tell the `APILogWorker` to send any currently enqueued logs and block until
         completion.
@@ -126,7 +136,7 @@ class APILogHandler(logging.Handler):
 
         return await APILogWorker.drain_all()
 
-    def emit(self, record: logging.LogRecord):
+    def emit(self, record: logging.LogRecord) -> None:
         """
         Send a log to the `APILogWorker`
         """
@@ -239,7 +249,7 @@ class APILogHandler(logging.Handler):
 
 
 class WorkerAPILogHandler(APILogHandler):
-    def emit(self, record: logging.LogRecord):
+    def emit(self, record: logging.LogRecord) -> None:
         # Open-source API servers do not currently support worker logs, and
         # worker logs only have an associated worker ID when connected to Cloud,
         # so we won't send worker logs to the API unless they have a worker ID.
@@ -278,13 +288,13 @@ class WorkerAPILogHandler(APILogHandler):
         return log
 
 
-class PrefectConsoleHandler(logging.StreamHandler):
+class PrefectConsoleHandler(StreamHandler):
     def __init__(
         self,
-        stream=None,
-        highlighter: Highlighter = PrefectConsoleHighlighter,
-        styles: Optional[Dict[str, str]] = None,
-        level: Union[int, str] = logging.NOTSET,
+        stream: TextIO | None = None,
+        highlighter: type[Highlighter] = PrefectConsoleHighlighter,
+        styles: dict[str, str] | None = None,
+        level: int | str = logging.NOTSET,
     ):
         """
         The default console handler for Prefect, which highlights log levels,
@@ -307,14 +317,14 @@ class PrefectConsoleHandler(logging.StreamHandler):
             theme = Theme(inherit=False)
 
         self.level = level
-        self.console = Console(
+        self.console: Console = Console(
             highlighter=highlighter,
             theme=theme,
             file=self.stream,
             markup=markup_console,
         )
 
-    def emit(self, record: logging.LogRecord):
+    def emit(self, record: logging.LogRecord) -> None:
         try:
             message = self.format(record)
             self.console.print(message, soft_wrap=True)

prefect/logging/highlighters.py

@@ -7,7 +7,7 @@ class LevelHighlighter(RegexHighlighter):
     """Apply style to log levels."""
 
     base_style = "level."
-    highlights = [
+    highlights: list[str] = [
         r"(?P<debug_level>DEBUG)",
         r"(?P<info_level>INFO)",
         r"(?P<warning_level>WARNING)",
@@ -20,7 +20,7 @@ class UrlHighlighter(RegexHighlighter):
     """Apply style to urls."""
 
     base_style = "url."
-    highlights = [
+    highlights: list[str] = [
         r"(?P<web_url>(https|http|ws|wss):\/\/[0-9a-zA-Z\$\-\_\+\!`\(\)\,\.\?\/\;\:\&\=\%\#]*)",
         r"(?P<local_url>(file):\/\/[0-9a-zA-Z\$\-\_\+\!`\(\)\,\.\?\/\;\:\&\=\%\#]*)",
     ]
@@ -30,7 +30,7 @@ class NameHighlighter(RegexHighlighter):
     """Apply style to names."""
 
     base_style = "name."
-    highlights = [
+    highlights: list[str] = [
         # ?i means case insensitive
         # ?<= means find string right after the words: flow run
         r"(?i)(?P<flow_run_name>(?<=flow run) \'(.*?)\')",
@@ -44,7 +44,7 @@ class StateHighlighter(RegexHighlighter):
     """Apply style to states."""
 
     base_style = "state."
-    highlights = [
+    highlights: list[str] = [
         rf"(?P<{state.lower()}_state>{state.title()})" for state in StateType
     ] + [
         r"(?P<cached_state>Cached)(?=\(type=COMPLETED\))"  # Highlight only "Cached"
@@ -55,7 +55,7 @@ class PrefectConsoleHighlighter(RegexHighlighter):
     """Applies style from multiple highlighters."""
 
     base_style = "log."
-    highlights = (
+    highlights: list[str] = (
         LevelHighlighter.highlights
         + UrlHighlighter.highlights
         + NameHighlighter.highlights