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

prefect/client/schemas/objects.py

@@ -1,15 +1,16 @@
 import datetime
 import warnings
+from collections.abc import Callable, Mapping
 from functools import partial
 from typing import (
     TYPE_CHECKING,
     Annotated,
     Any,
-    Dict,
+    ClassVar,
     Generic,
-    List,
     Optional,
     Union,
+    cast,
     overload,
 )
 from uuid import UUID, uuid4
@@ -23,13 +24,12 @@ from pydantic import (
     HttpUrl,
     IPvAnyNetwork,
     SerializationInfo,
+    SerializerFunctionWrapHandler,
     Tag,
     field_validator,
     model_serializer,
     model_validator,
 )
-from pydantic.functional_validators import ModelWrapValidatorHandler
-from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import Literal, Self, TypeVar
 
 from prefect._internal.compatibility import deprecated
@@ -64,8 +64,13 @@ from prefect.utilities.names import generate_slug
 from prefect.utilities.pydantic import handle_secret_render
 
 if TYPE_CHECKING:
+    from prefect.client.schemas.actions import StateCreate
     from prefect.results import BaseResult, ResultRecordMetadata
 
+    DateTime = pendulum.DateTime
+else:
+    from pydantic_extra_types.pendulum_dt import DateTime
+
 
 R = TypeVar("R", default=Any)
 
@@ -180,7 +185,7 @@ class StateDetails(PrefectBaseModel):
     pause_timeout: Optional[DateTime] = None
     pause_reschedule: bool = False
     pause_key: Optional[str] = None
-    run_input_keyset: Optional[Dict[str, str]] = None
+    run_input_keyset: Optional[dict[str, str]] = None
     refresh_cache: Optional[bool] = None
     retriable: Optional[bool] = None
     transition_id: Optional[UUID] = None
@@ -215,11 +220,21 @@ class State(ObjectBaseModel, Generic[R]):
     ] = Field(default=None)
 
     @overload
-    def result(self: "State[R]", raise_on_failure: bool = True) -> R:
+    def result(
+        self: "State[R]",
+        raise_on_failure: Literal[True] = ...,
+        fetch: bool = ...,
+        retry_result_failure: bool = ...,
+    ) -> R:
         ...
 
     @overload
-    def result(self: "State[R]", raise_on_failure: bool = False) -> Union[R, Exception]:
+    def result(
+        self: "State[R]",
+        raise_on_failure: Literal[False] = False,
+        fetch: bool = ...,
+        retry_result_failure: bool = ...,
+    ) -> Union[R, Exception]:
         ...
 
     @deprecated.deprecated_parameter(
@@ -311,7 +326,7 @@ class State(ObjectBaseModel, Generic[R]):
             retry_result_failure=retry_result_failure,
         )
 
-    def to_state_create(self):
+    def to_state_create(self) -> "StateCreate":
         """
         Convert this state to a `StateCreate` type which can be used to set the state of
         a run in the API.
@@ -327,7 +342,7 @@ class State(ObjectBaseModel, Generic[R]):
         )
 
         if isinstance(self.data, BaseResult):
-            data = self.data
+            data = cast(BaseResult[R], self.data)
         elif isinstance(self.data, ResultRecord) and should_persist_result():
             data = self.data.metadata
         else:
@@ -348,14 +363,14 @@ class State(ObjectBaseModel, Generic[R]):
         # validation check and an error will be raised after this function is called
         name = self.name
         if name is None and self.type:
-            self.name = " ".join([v.capitalize() for v in self.type.value.split("_")])
+            self.name = " ".join([v.capitalize() for v in self.type.split("_")])
         return self
 
     @model_validator(mode="after")
     def default_scheduled_start_time(self) -> Self:
         if self.type == StateType.SCHEDULED:
             if not self.state_details.scheduled_time:
-                self.state_details.scheduled_time = DateTime.now("utc")
+                self.state_details.scheduled_time = pendulum.DateTime.now("utc")
         return self
 
     @model_validator(mode="after")
@@ -395,17 +410,19 @@ class State(ObjectBaseModel, Generic[R]):
         return self.type == StateType.PAUSED
 
     def model_copy(
-        self, *, update: Optional[Dict[str, Any]] = None, deep: bool = False
-    ):
+        self, *, update: Optional[Mapping[str, Any]] = None, deep: bool = False
+    ) -> Self:
         """
         Copying API models should return an object that could be inserted into the
         database again. The 'timestamp' is reset using the default factory.
         """
-        update = update or {}
-        update.setdefault("timestamp", self.model_fields["timestamp"].get_default())
+        update = {
+            "timestamp": self.model_fields["timestamp"].get_default(),
+            **(update or {}),
+        }
         return super().model_copy(update=update, deep=deep)
 
-    def fresh_copy(self, **kwargs) -> Self:
+    def fresh_copy(self, **kwargs: Any) -> Self:
         """
         Return a fresh copy of the state with a new ID.
         """
@@ -443,12 +460,14 @@ class State(ObjectBaseModel, Generic[R]):
         `MyCompletedState("my message", type=COMPLETED)`
         """
 
-        display = []
+        display: list[str] = []
 
         if self.message:
             display.append(repr(self.message))
 
-        if self.type.value.lower() != self.name.lower():
+        if TYPE_CHECKING:
+            assert self.name is not None
+        if self.type.lower() != self.name.lower():
             display.append(f"type={self.type.value}")
 
         return f"{self.name}({', '.join(display)})"
@@ -487,7 +506,7 @@ class FlowRunPolicy(PrefectBaseModel):
     retry_delay: Optional[int] = Field(
         default=None, description="The delay time between retries, in seconds."
     )
-    pause_keys: Optional[set] = Field(
+    pause_keys: Optional[set[str]] = Field(
         default_factory=set, description="Tracks pauses this run has observed."
     )
     resuming: Optional[bool] = Field(
@@ -499,7 +518,7 @@ class FlowRunPolicy(PrefectBaseModel):
 
     @model_validator(mode="before")
     @classmethod
-    def populate_deprecated_fields(cls, values: Any):
+    def populate_deprecated_fields(cls, values: Any) -> Any:
         if isinstance(values, dict):
             return set_run_policy_deprecated_fields(values)
         return values
@@ -536,7 +555,7 @@ class FlowRun(ObjectBaseModel):
         description="The version of the flow executed in this flow run.",
         examples=["1.0"],
     )
-    parameters: Dict[str, Any] = Field(
+    parameters: dict[str, Any] = Field(
         default_factory=dict, description="Parameters for the flow run."
     )
     idempotency_key: Optional[str] = Field(
@@ -546,7 +565,7 @@ class FlowRun(ObjectBaseModel):
             " run is not created multiple times."
         ),
     )
-    context: Dict[str, Any] = Field(
+    context: dict[str, Any] = Field(
         default_factory=dict,
         description="Additional context for the flow run.",
         examples=[{"my_var": "my_val"}],
@@ -554,12 +573,12 @@ class FlowRun(ObjectBaseModel):
     empirical_policy: FlowRunPolicy = Field(
         default_factory=FlowRunPolicy,
     )
-    tags: List[str] = Field(
+    tags: list[str] = Field(
         default_factory=list,
         description="A list of tags on the flow run",
         examples=[["tag-1", "tag-2"]],
     )
-    labels: KeyValueLabelsField
+    labels: KeyValueLabelsField = Field(default_factory=dict)
     parent_task_run_id: Optional[UUID] = Field(
         default=None,
         description=(
@@ -632,7 +651,7 @@ class FlowRun(ObjectBaseModel):
         description="The state of the flow run.",
         examples=["State(type=StateType.COMPLETED)"],
     )
-    job_variables: Optional[dict] = Field(
+    job_variables: Optional[dict[str, Any]] = Field(
         default=None,
         description="Job variables for the flow run.",
     )
@@ -663,7 +682,7 @@ class FlowRun(ObjectBaseModel):
 
     @field_validator("name", mode="before")
     @classmethod
-    def set_default_name(cls, name):
+    def set_default_name(cls, name: Optional[str]) -> str:
         return get_or_create_run_name(name)
 
 
@@ -687,7 +706,7 @@ class TaskRunPolicy(PrefectBaseModel):
         deprecated=True,
     )
     retries: Optional[int] = Field(default=None, description="The number of retries.")
-    retry_delay: Union[None, int, List[int]] = Field(
+    retry_delay: Union[None, int, list[int]] = Field(
         default=None,
         description="A delay time or list of delay times between retries, in seconds.",
     )
@@ -710,18 +729,20 @@ class TaskRunPolicy(PrefectBaseModel):
                 self.retries = self.max_retries
 
             if not self.retry_delay and self.retry_delay_seconds != 0:
-                self.retry_delay = self.retry_delay_seconds
+                self.retry_delay = int(self.retry_delay_seconds)
 
         return self
 
     @field_validator("retry_delay")
     @classmethod
-    def validate_configured_retry_delays(cls, v):
+    def validate_configured_retry_delays(
+        cls, v: Optional[list[float]]
+    ) -> Optional[list[float]]:
         return list_length_50_or_less(v)
 
     @field_validator("retry_jitter_factor")
     @classmethod
-    def validate_jitter_factor(cls, v):
+    def validate_jitter_factor(cls, v: Optional[float]) -> Optional[float]:
         return validate_not_negative(v)
 
 
@@ -731,9 +752,11 @@ class TaskRunInput(PrefectBaseModel):
     could include, constants, parameters, or other task runs.
     """
 
-    model_config = ConfigDict(frozen=True)
+    model_config: ClassVar[ConfigDict] = ConfigDict(frozen=True)
 
-    input_type: str
+    if not TYPE_CHECKING:
+        # subclasses provide the concrete type for this field
+        input_type: str
 
 
 class TaskRunResult(TaskRunInput):
@@ -791,16 +814,16 @@ class TaskRun(ObjectBaseModel):
     empirical_policy: TaskRunPolicy = Field(
         default_factory=TaskRunPolicy,
     )
-    tags: List[str] = Field(
+    tags: list[str] = Field(
         default_factory=list,
         description="A list of tags for the task run.",
         examples=[["tag-1", "tag-2"]],
     )
-    labels: KeyValueLabelsField
+    labels: KeyValueLabelsField = Field(default_factory=dict)
     state_id: Optional[UUID] = Field(
         default=None, description="The id of the current task run state."
     )
-    task_inputs: Dict[str, List[Union[TaskRunResult, Parameter, Constant]]] = Field(
+    task_inputs: dict[str, list[Union[TaskRunResult, Parameter, Constant]]] = Field(
         default_factory=dict,
         description=(
             "Tracks the source of inputs to a task run. Used for internal bookkeeping. "
@@ -865,7 +888,7 @@ class TaskRun(ObjectBaseModel):
 
     @field_validator("name", mode="before")
     @classmethod
-    def set_default_name(cls, name):
+    def set_default_name(cls, name: Optional[str]) -> Name:
         return get_or_create_run_name(name)
 
 
@@ -883,7 +906,7 @@ class Workspace(PrefectBaseModel):
     workspace_name: str = Field(..., description="The workspace name.")
     workspace_description: str = Field(..., description="Description of the workspace.")
     workspace_handle: str = Field(..., description="The workspace's unique handle.")
-    model_config = ConfigDict(extra="ignore")
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="ignore")
 
     @property
     def handle(self) -> str:
@@ -912,7 +935,7 @@ class Workspace(PrefectBaseModel):
             f"/workspace/{self.workspace_id}"
         )
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash(self.handle)
 
 
@@ -935,7 +958,7 @@ class IPAllowlist(PrefectBaseModel):
     Expected payload for an IP allowlist from the Prefect Cloud API.
     """
 
-    entries: List[IPAllowlistEntry]
+    entries: list[IPAllowlistEntry]
 
 
 class IPAllowlistMyAccessResponse(PrefectBaseModel):
@@ -973,14 +996,14 @@ class BlockSchema(ObjectBaseModel):
     """A representation of a block schema."""
 
     checksum: str = Field(default=..., description="The block schema's unique checksum")
-    fields: Dict[str, Any] = Field(
+    fields: dict[str, Any] = Field(
         default_factory=dict, description="The block schema's field schema"
     )
     block_type_id: Optional[UUID] = Field(default=..., description="A block type ID")
     block_type: Optional[BlockType] = Field(
         default=None, description="The associated block type"
     )
-    capabilities: List[str] = Field(
+    capabilities: list[str] = Field(
         default_factory=list,
         description="A list of Block capabilities",
     )
@@ -999,7 +1022,7 @@ class BlockDocument(ObjectBaseModel):
             "The block document's name. Not required for anonymous block documents."
         ),
     )
-    data: Dict[str, Any] = Field(
+    data: dict[str, Any] = Field(
         default_factory=dict, description="The block document's data"
     )
     block_schema_id: UUID = Field(default=..., description="A block schema ID")
@@ -1011,7 +1034,7 @@ class BlockDocument(ObjectBaseModel):
     block_type: Optional[BlockType] = Field(
         default=None, description="The associated block type"
     )
-    block_document_references: Dict[str, Dict[str, Any]] = Field(
+    block_document_references: dict[str, dict[str, Any]] = Field(
         default_factory=dict, description="Record of the block document's references"
     )
     is_anonymous: bool = Field(
@@ -1026,13 +1049,15 @@ class BlockDocument(ObjectBaseModel):
 
     @model_validator(mode="before")
     @classmethod
-    def validate_name_is_present_if_not_anonymous(cls, values):
+    def validate_name_is_present_if_not_anonymous(
+        cls, values: dict[str, Any]
+    ) -> dict[str, Any]:
         return validate_name_present_on_nonanonymous_blocks(values)
 
     @model_serializer(mode="wrap")
     def serialize_data(
-        self, handler: ModelWrapValidatorHandler, info: SerializationInfo
-    ):
+        self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
+    ) -> Any:
         self.data = visit_collection(
             self.data,
             visit_fn=partial(handle_secret_render, context=info.context or {}),
@@ -1047,7 +1072,7 @@ class Flow(ObjectBaseModel):
     name: Name = Field(
         default=..., description="The name of the flow", examples=["my-flow"]
     )
-    tags: List[str] = Field(
+    tags: list[str] = Field(
         default_factory=list,
         description="A list of flow tags",
         examples=[["tag-1", "tag-2"]],
@@ -1091,22 +1116,22 @@ class Deployment(ObjectBaseModel):
     concurrency_limit: Optional[int] = Field(
         default=None, description="The concurrency limit for the deployment."
     )
-    schedules: List[DeploymentSchedule] = Field(
+    schedules: list[DeploymentSchedule] = Field(
         default_factory=list, description="A list of schedules for the deployment."
     )
-    job_variables: Dict[str, Any] = Field(
+    job_variables: dict[str, Any] = Field(
         default_factory=dict,
         description="Overrides to apply to flow run infrastructure at runtime.",
     )
-    parameters: Dict[str, Any] = Field(
+    parameters: dict[str, Any] = Field(
         default_factory=dict,
         description="Parameters for flow runs scheduled by the deployment.",
     )
-    pull_steps: Optional[List[dict]] = Field(
+    pull_steps: Optional[list[dict[str, Any]]] = Field(
         default=None,
         description="Pull steps for cloning and running this deployment.",
     )
-    tags: List[str] = Field(
+    tags: list[str] = Field(
         default_factory=list,
         description="A list of tags for the deployment",
         examples=[["tag-1", "tag-2"]],
@@ -1123,7 +1148,7 @@ class Deployment(ObjectBaseModel):
         default=None,
         description="The last time the deployment was polled for status updates.",
     )
-    parameter_openapi_schema: Optional[Dict[str, Any]] = Field(
+    parameter_openapi_schema: Optional[dict[str, Any]] = Field(
         default=None,
         description="The parameter schema of the flow, including defaults.",
     )
@@ -1177,7 +1202,7 @@ class ConcurrencyLimit(ObjectBaseModel):
         default=..., description="A tag the concurrency limit is applied to."
     )
     concurrency_limit: int = Field(default=..., description="The concurrency limit.")
-    active_slots: List[UUID] = Field(
+    active_slots: list[UUID] = Field(
         default_factory=list,
         description="A list of active run ids using a concurrency slot",
     )
@@ -1224,7 +1249,7 @@ class BlockDocumentReference(ObjectBaseModel):
 
     @model_validator(mode="before")
     @classmethod
-    def validate_parent_and_ref_are_different(cls, values):
+    def validate_parent_and_ref_are_different(cls, values: Any) -> Any:
         if isinstance(values, dict):
             return validate_parent_and_ref_diff(values)
         return values
@@ -1234,7 +1259,7 @@ class Configuration(ObjectBaseModel):
     """An ORM representation of account info."""
 
     key: str = Field(default=..., description="Account info key")
-    value: Dict[str, Any] = Field(default=..., description="Account info")
+    value: dict[str, Any] = Field(default=..., description="Account info")
 
 
 class SavedSearchFilter(PrefectBaseModel):
@@ -1258,7 +1283,7 @@ class SavedSearch(ObjectBaseModel):
     """An ORM representation of saved search data. Represents a set of filter criteria."""
 
     name: str = Field(default=..., description="The name of the saved search.")
-    filters: List[SavedSearchFilter] = Field(
+    filters: list[SavedSearchFilter] = Field(
         default_factory=list, description="The filter set for the saved search."
     )
 
@@ -1281,11 +1306,11 @@ class Log(ObjectBaseModel):
 class QueueFilter(PrefectBaseModel):
     """Filter criteria definition for a work queue."""
 
-    tags: Optional[List[str]] = Field(
+    tags: Optional[list[str]] = Field(
         default=None,
         description="Only include flow runs with these tags in the work queue.",
     )
-    deployment_ids: Optional[List[UUID]] = Field(
+    deployment_ids: Optional[list[UUID]] = Field(
         default=None,
         description="Only include flow runs from these deployments in the work queue.",
     )
@@ -1345,7 +1370,7 @@ class WorkQueueHealthPolicy(PrefectBaseModel):
     )
 
     def evaluate_health_status(
-        self, late_runs_count: int, last_polled: Optional[DateTime] = None
+        self, late_runs_count: int, last_polled: Optional[pendulum.DateTime] = None
     ) -> bool:
         """
         Given empirical information about the state of the work queue, evaluate its health status.
@@ -1397,10 +1422,10 @@ class FlowRunNotificationPolicy(ObjectBaseModel):
     is_active: bool = Field(
         default=True, description="Whether the policy is currently active"
     )
-    state_names: List[str] = Field(
+    state_names: list[str] = Field(
         default=..., description="The flow run states that trigger notifications"
     )
-    tags: List[str] = Field(
+    tags: list[str] = Field(
         default=...,
         description="The flow run tags that trigger notifications (set [] to disable)",
     )
@@ -1422,7 +1447,7 @@ class FlowRunNotificationPolicy(ObjectBaseModel):
 
     @field_validator("message_template")
     @classmethod
-    def validate_message_template_variables(cls, v):
+    def validate_message_template_variables(cls, v: Optional[str]) -> Optional[str]:
         return validate_message_template_variables(v)
 
 
@@ -1454,7 +1479,7 @@ class WorkPool(ObjectBaseModel):
         default=None, description="A description of the work pool."
     )
     type: str = Field(description="The work pool type.")
-    base_job_template: Dict[str, Any] = Field(
+    base_job_template: dict[str, Any] = Field(
         default_factory=dict, description="The work pool's base job template."
     )
     is_paused: bool = Field(
@@ -1469,10 +1494,12 @@ class WorkPool(ObjectBaseModel):
     )
 
     # this required field has a default of None so that the custom validator
-    # below will be called and produce a more helpful error message
-    default_queue_id: UUID = Field(
-        None, description="The id of the pool's default queue."
-    )
+    # below will be called and produce a more helpful error message. Because
+    # the field metadata is attached via an annotation, the default is hidden
+    # from type checkers.
+    default_queue_id: Annotated[
+        UUID, Field(default=None, description="The id of the pool's default queue.")
+    ]
 
     @property
     def is_push_pool(self) -> bool:
@@ -1484,7 +1511,7 @@ class WorkPool(ObjectBaseModel):
 
     @field_validator("default_queue_id")
     @classmethod
-    def helpful_error_for_missing_default_queue_id(cls, v):
+    def helpful_error_for_missing_default_queue_id(cls, v: Optional[UUID]) -> UUID:
         return validate_default_queue_id_not_none(v)
 
 
@@ -1495,8 +1522,8 @@ class Worker(ObjectBaseModel):
     work_pool_id: UUID = Field(
         description="The work pool with which the queue is associated."
     )
-    last_heartbeat_time: datetime.datetime = Field(
-        None, description="The last time the worker process sent a heartbeat."
+    last_heartbeat_time: Optional[datetime.datetime] = Field(
+        default=None, description="The last time the worker process sent a heartbeat."
     )
     heartbeat_interval_seconds: Optional[int] = Field(
         default=None,
@@ -1529,14 +1556,14 @@ class Artifact(ObjectBaseModel):
         default=None, description="A markdown-enabled description of the artifact."
     )
     # data will eventually be typed as `Optional[Union[Result, Any]]`
-    data: Optional[Union[Dict[str, Any], Any]] = Field(
+    data: Optional[Union[dict[str, Any], Any]] = Field(
         default=None,
         description=(
             "Data associated with the artifact, e.g. a result.; structure depends on"
             " the artifact type."
         ),
     )
-    metadata_: Optional[Dict[str, str]] = Field(
+    metadata_: Optional[dict[str, str]] = Field(
         default=None,
         description=(
             "User-defined artifact metadata. Content must be string key and value"
@@ -1552,7 +1579,9 @@ class Artifact(ObjectBaseModel):
 
     @field_validator("metadata_")
     @classmethod
-    def validate_metadata_length(cls, v):
+    def validate_metadata_length(
+        cls, v: Optional[dict[str, str]]
+    ) -> Optional[dict[str, str]]:
         return validate_max_metadata_length(v)
 
 
@@ -1571,14 +1600,14 @@ class ArtifactCollection(ObjectBaseModel):
     description: Optional[str] = Field(
         default=None, description="A markdown-enabled description of the artifact."
     )
-    data: Optional[Union[Dict[str, Any], Any]] = Field(
+    data: Optional[Union[dict[str, Any], Any]] = Field(
         default=None,
         description=(
             "Data associated with the artifact, e.g. a result.; structure depends on"
             " the artifact type."
         ),
     )
-    metadata_: Optional[Dict[str, str]] = Field(
+    metadata_: Optional[dict[str, str]] = Field(
         default=None,
         description=(
             "User-defined artifact metadata. Content must be string key and value"
@@ -1605,7 +1634,7 @@ class Variable(ObjectBaseModel):
         description="The value of the variable",
         examples=["my_value"],
     )
-    tags: List[str] = Field(
+    tags: list[str] = Field(
         default_factory=list,
         description="A list of variable tags",
         examples=[["tag-1", "tag-2"]],
@@ -1630,7 +1659,7 @@ class FlowRunInput(ObjectBaseModel):
 
     @field_validator("key", check_fields=False)
     @classmethod
-    def validate_name_characters(cls, v):
+    def validate_name_characters(cls, v: str) -> str:
         raise_on_name_alphanumeric_dashes_only(v)
         return v
 
@@ -1675,7 +1704,7 @@ class CsrfToken(ObjectBaseModel):
     )
 
 
-__getattr__ = getattr_migration(__name__)
+__getattr__: Callable[[str], Any] = getattr_migration(__name__)
 
 
 class Integration(PrefectBaseModel):
@@ -1693,7 +1722,7 @@ class WorkerMetadata(PrefectBaseModel):
     should support flexible metadata.
     """
 
-    integrations: List[Integration] = Field(
+    integrations: list[Integration] = Field(
         default=..., description="Prefect integrations installed in the worker."
     )
-    model_config = ConfigDict(extra="allow")
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")