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

prefect/engine.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 from uuid import UUID
 
 from prefect._internal.compatibility.migration import getattr_migration
@@ -15,12 +15,19 @@ from prefect.utilities.asyncutils import (
     run_coro_as_sync,
 )
 
-engine_logger = get_logger("engine")
+if TYPE_CHECKING:
+    import logging
+
+    from prefect.flow_engine import FlowRun
+    from prefect.flows import Flow
+    from prefect.logging.loggers import LoggingAdapter
+
+engine_logger: "logging.Logger" = get_logger("engine")
 
 
 if __name__ == "__main__":
     try:
-        flow_run_id = UUID(
+        flow_run_id: UUID = UUID(
             sys.argv[1] if len(sys.argv) > 1 else os.environ.get("PREFECT__FLOW_RUN_ID")
         )
     except Exception:
@@ -37,11 +44,11 @@ if __name__ == "__main__":
             run_flow,
         )
 
-        flow_run = load_flow_run(flow_run_id=flow_run_id)
-        run_logger = flow_run_logger(flow_run=flow_run)
+        flow_run: "FlowRun" = load_flow_run(flow_run_id=flow_run_id)
+        run_logger: "LoggingAdapter" = flow_run_logger(flow_run=flow_run)
 
         try:
-            flow = load_flow(flow_run)
+            flow: "Flow[..., Any]" = load_flow(flow_run)
         except Exception:
             run_logger.error(
                 "Unexpected exception encountered when trying to load flow",
@@ -55,15 +62,17 @@ if __name__ == "__main__":
         else:
             run_flow(flow, flow_run=flow_run, error_logger=run_logger)
 
-    except Abort as exc:
+    except Abort as abort_signal:
+        abort_signal: Abort
         engine_logger.info(
             f"Engine execution of flow run '{flow_run_id}' aborted by orchestrator:"
-            f" {exc}"
+            f" {abort_signal}"
         )
         exit(0)
-    except Pause as exc:
+    except Pause as pause_signal:
+        pause_signal: Pause
         engine_logger.info(
-            f"Engine execution of flow run '{flow_run_id}' is paused: {exc}"
+            f"Engine execution of flow run '{flow_run_id}' is paused: {pause_signal}"
         )
         exit(0)
     except Exception:

prefect/events/cli/automations.py

@@ -3,7 +3,7 @@ Command line interface for working with automations.
 """
 
 import functools
-from typing import Optional, Type
+from typing import Any, Callable, Optional, Type
 from uuid import UUID
 
 import orjson
@@ -21,16 +21,16 @@ from prefect.client.orchestration import get_client
 from prefect.events.schemas.automations import Automation
 from prefect.exceptions import PrefectHTTPStatusError
 
-automations_app = PrefectTyper(
+automations_app: PrefectTyper = PrefectTyper(
     name="automation",
     help="Manage automations.",
 )
 app.add_typer(automations_app, aliases=["automations"])
 
 
-def requires_automations(func):
+def requires_automations(func: Callable[..., Any]) -> Callable[..., Any]:
     @functools.wraps(func)
-    async def wrapper(*args, **kwargs):
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
         try:
             return await func(*args, **kwargs)
         except RuntimeError as exc:

prefect/events/clients.py

@@ -70,18 +70,21 @@ EVENT_WEBSOCKET_CHECKPOINTS = Counter(
     labelnames=["client"],
 )
 
-logger = get_logger(__name__)
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger(__name__)
 
 
-def http_to_ws(url: str):
+def http_to_ws(url: str) -> str:
     return url.replace("https://", "wss://").replace("http://", "ws://").rstrip("/")
 
 
-def events_in_socket_from_api_url(url: str):
+def events_in_socket_from_api_url(url: str) -> str:
     return http_to_ws(url) + "/events/in"
 
 
-def events_out_socket_from_api_url(url: str):
+def events_out_socket_from_api_url(url: str) -> str:
     return http_to_ws(url) + "/events/out"
 
 
@@ -250,11 +253,11 @@ class AssertingEventsClient(EventsClient):
     last: ClassVar["Optional[AssertingEventsClient]"] = None
     all: ClassVar[List["AssertingEventsClient"]] = []
 
-    args: Tuple
-    kwargs: Dict[str, Any]
-    events: List[Event]
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
+    events: list[Event]
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any):
         AssertingEventsClient.last = self
         AssertingEventsClient.all.append(self)
         self.args = args
@@ -431,13 +434,13 @@ class AssertingPassthroughEventsClient(PrefectEventsClient):
     during tests AND sends them to a Prefect server."""
 
     last: ClassVar["Optional[AssertingPassthroughEventsClient]"] = None
-    all: ClassVar[List["AssertingPassthroughEventsClient"]] = []
+    all: ClassVar[list["AssertingPassthroughEventsClient"]] = []
 
-    args: Tuple
-    kwargs: Dict[str, Any]
-    events: List[Event]
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
+    events: list[Event]
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
         AssertingPassthroughEventsClient.last = self
         AssertingPassthroughEventsClient.all.append(self)
@@ -449,7 +452,7 @@ class AssertingPassthroughEventsClient(PrefectEventsClient):
         cls.last = None
         cls.all = []
 
-    def pop_events(self) -> List[Event]:
+    def pop_events(self) -> list[Event]:
         events = self.events
         self.events = []
         return events

prefect/events/schemas/automations.py

@@ -52,7 +52,7 @@ class Trigger(PrefectBaseModel, abc.ABC, extra="ignore"):  # type: ignore[call-a
 
     _deployment_id: Optional[UUID] = PrivateAttr(default=None)
 
-    def set_deployment_id(self, deployment_id: UUID):
+    def set_deployment_id(self, deployment_id: UUID) -> None:
         self._deployment_id = deployment_id
 
     def owner_resource(self) -> Optional[str]:
@@ -277,7 +277,7 @@ class MetricTriggerQuery(PrefectBaseModel):
     )
 
     @field_validator("range", "firing_for")
-    def enforce_minimum_range(cls, value: timedelta):
+    def enforce_minimum_range(cls, value: timedelta) -> timedelta:
         if value < timedelta(seconds=300):
             raise ValueError("The minimum range is 300 seconds (5 minutes)")
         return value
@@ -404,13 +404,17 @@ class AutomationCore(PrefectBaseModel, extra="ignore"):  # type: ignore[call-arg
     """Defines an action a user wants to take when a certain number of events
     do or don't happen to the matching resources"""
 
-    name: str = Field(..., description="The name of this automation")
-    description: str = Field("", description="A longer description of this automation")
+    name: str = Field(default=..., description="The name of this automation")
+    description: str = Field(
+        default="", description="A longer description of this automation"
+    )
 
-    enabled: bool = Field(True, description="Whether this automation will be evaluated")
+    enabled: bool = Field(
+        default=True, description="Whether this automation will be evaluated"
+    )
 
     trigger: TriggerTypes = Field(
-        ...,
+        default=...,
         description=(
             "The criteria for which events this Automation covers and how it will "
             "respond to the presence or absence of those events"
@@ -418,7 +422,7 @@ class AutomationCore(PrefectBaseModel, extra="ignore"):  # type: ignore[call-arg
     )
 
     actions: List[ActionTypes] = Field(
-        ...,
+        default=...,
         description="The actions to perform when this Automation triggers",
     )
 
@@ -438,4 +442,4 @@ class AutomationCore(PrefectBaseModel, extra="ignore"):  # type: ignore[call-arg
 
 
 class Automation(AutomationCore):
-    id: UUID = Field(..., description="The ID of this automation")
+    id: UUID = Field(default=..., description="The ID of this automation")

prefect/events/schemas/events.py

@@ -1,6 +1,7 @@
 import copy
 from collections import defaultdict
 from typing import (
+    TYPE_CHECKING,
     Any,
     ClassVar,
     Dict,
@@ -32,7 +33,10 @@ from prefect.types import DateTime
 
 from .labelling import Labelled
 
-logger = get_logger(__name__)
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger(__name__)
 
 
 class Resource(Labelled):

prefect/events/worker.py

@@ -82,7 +82,7 @@ class EventsWorker(QueueService[Event]):
 
         await self._client.emit(event)
 
-    async def attach_related_resources_from_context(self, event: Event):
+    async def attach_related_resources_from_context(self, event: Event) -> None:
         if "prefect.resource.lineage-group" in event.resource:
             # We attach related resources to lineage events in `emit_lineage_event`,
             # instead of the worker, because not all run-related resources are

prefect/filesystems.py

@@ -281,7 +281,7 @@ class RemoteFileSystem(WritableFileSystem, WritableDeploymentStorage):
     _filesystem: fsspec.AbstractFileSystem = None
 
     @field_validator("basepath")
-    def check_basepath(cls, value):
+    def check_basepath(cls, value: str) -> str:
         return validate_basepath(value)
 
     def _resolve_path(self, path: str) -> str:

prefect/flow_engine.py

@@ -146,7 +146,7 @@ class BaseFlowRunEngine(Generic[P, R]):
     _flow_run_name_set: bool = False
     _telemetry: RunTelemetry = field(default_factory=RunTelemetry)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.flow is None and self.flow_run_id is None:
             raise ValueError("Either a flow or a flow_run_id must be provided.")
 
@@ -167,7 +167,7 @@ class BaseFlowRunEngine(Generic[P, R]):
             return False  # TODO: handle this differently?
         return getattr(self, "flow_run").state.is_pending()
 
-    def cancel_all_tasks(self):
+    def cancel_all_tasks(self) -> None:
         if hasattr(self.flow.task_runner, "cancel_all"):
             self.flow.task_runner.cancel_all()  # type: ignore
 
@@ -208,6 +208,8 @@ class BaseFlowRunEngine(Generic[P, R]):
 @dataclass
 class FlowRunEngine(BaseFlowRunEngine[P, R]):
     _client: Optional[SyncPrefectClient] = None
+    flow_run: FlowRun | None = None
+    parameters: dict[str, Any] | None = None
 
     @property
     def client(self) -> SyncPrefectClient:
@@ -502,7 +504,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
             tags=TagsContext.get().current_tags,
         )
 
-    def call_hooks(self, state: Optional[State] = None):
+    def call_hooks(self, state: Optional[State] = None) -> None:
         if state is None:
             state = self.state
         flow = self.flow
@@ -600,7 +602,9 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
 
             # set the logger to the flow run logger
 
-            self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
+            self.logger: "logging.Logger" = flow_run_logger(
+                flow_run=self.flow_run, flow=self.flow
+            )  # type: ignore
 
             # update the flow run name if necessary
             if not self._flow_run_name_set and self.flow.flow_run_name:
@@ -768,6 +772,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
     """
 
     _client: Optional[PrefectClient] = None
+    parameters: dict[str, Any] | None = None
+    flow_run: FlowRun | None = None
 
     @property
     def client(self) -> PrefectClient:
@@ -1061,7 +1067,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             tags=TagsContext.get().current_tags,
         )
 
-    async def call_hooks(self, state: Optional[State] = None):
+    async def call_hooks(self, state: Optional[State] = None) -> None:
         if state is None:
             state = self.state
         flow = self.flow
@@ -1158,7 +1164,9 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             stack.enter_context(ConcurrencyContext())
 
             # set the logger to the flow run logger
-            self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
+            self.logger: "logging.Logger" = flow_run_logger(
+                flow_run=self.flow_run, flow=self.flow
+            )
 
             # update the flow run name if necessary
 
@@ -1320,7 +1328,7 @@ def run_flow_sync(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
     return_type: Literal["state", "result"] = "result",
 ) -> Union[R, State, None]:
     engine = FlowRunEngine[P, R](
@@ -1342,7 +1350,7 @@ async def run_flow_async(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
     return_type: Literal["state", "result"] = "result",
 ) -> Union[R, State, None]:
     engine = AsyncFlowRunEngine[P, R](
@@ -1361,7 +1369,7 @@ def run_generator_flow_sync(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
     return_type: Literal["state", "result"] = "result",
 ) -> Generator[R, None, None]:
     if return_type != "result":