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

prefect/client/schemas/responses.py

@@ -1,5 +1,5 @@
 import datetime
-from typing import Any, Dict, List, Optional, TypeVar, Union
+from typing import Any, ClassVar, Generic, Optional, TypeVar, Union
 from uuid import UUID
 
 from pydantic import ConfigDict, Field
@@ -9,10 +9,11 @@ from typing_extensions import Literal
 import prefect.client.schemas.objects as objects
 from prefect._internal.schemas.bases import ObjectBaseModel, PrefectBaseModel
 from prefect._internal.schemas.fields import CreatedBy, UpdatedBy
+from prefect.types import KeyValueLabelsField
 from prefect.utilities.collections import AutoEnum
 from prefect.utilities.names import generate_slug
 
-R = TypeVar("R")
+T = TypeVar("T")
 
 
 class SetStateStatus(AutoEnum):
@@ -119,7 +120,7 @@ class HistoryResponse(PrefectBaseModel):
     interval_end: DateTime = Field(
         default=..., description="The end date of the interval."
     )
-    states: List[HistoryResponseState] = Field(
+    states: list[HistoryResponseState] = Field(
         default=..., description="A list of state histories during the interval."
     )
 
@@ -129,18 +130,18 @@ StateResponseDetails = Union[
 ]
 
 
-class OrchestrationResult(PrefectBaseModel):
+class OrchestrationResult(PrefectBaseModel, Generic[T]):
     """
     A container for the output of state orchestration.
     """
 
-    state: Optional[objects.State]
+    state: Optional[objects.State[T]]
     status: SetStateStatus
     details: StateResponseDetails
 
 
 class WorkerFlowRunResponse(PrefectBaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+    model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
 
     work_pool_id: UUID
     work_queue_id: UUID
@@ -178,7 +179,7 @@ class FlowRunResponse(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(
@@ -188,7 +189,7 @@ class FlowRunResponse(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"}],
@@ -196,11 +197,12 @@ class FlowRunResponse(ObjectBaseModel):
     empirical_policy: objects.FlowRunPolicy = Field(
         default_factory=objects.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
     parent_task_run_id: Optional[UUID] = Field(
         default=None,
         description=(
@@ -273,7 +275,7 @@ class FlowRunResponse(ObjectBaseModel):
         description="The state of the flow run.",
         examples=["objects.State(type=objects.StateType.COMPLETED)"],
     )
-    job_variables: Optional[dict] = Field(
+    job_variables: Optional[dict[str, Any]] = Field(
         default=None, description="Job variables for the flow run."
     )
 
@@ -333,26 +335,27 @@ class DeploymentResponse(ObjectBaseModel):
         default=None,
         description="The concurrency options for the deployment.",
     )
-    schedules: List[objects.DeploymentSchedule] = Field(
+    schedules: list[objects.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"]],
     )
+    labels: KeyValueLabelsField
     work_queue_name: Optional[str] = Field(
         default=None,
         description=(
@@ -364,7 +367,7 @@ class DeploymentResponse(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.",
     )
@@ -397,7 +400,7 @@ class DeploymentResponse(ObjectBaseModel):
         default=None,
         description="Optional information about the updater of this deployment.",
     )
-    work_queue_id: UUID = Field(
+    work_queue_id: Optional[UUID] = Field(
         default=None,
         description=(
             "The id of the work pool queue to which this deployment is assigned."
@@ -420,7 +423,7 @@ class DeploymentResponse(ObjectBaseModel):
 
 
 class MinimalConcurrencyLimitResponse(PrefectBaseModel):
-    model_config = ConfigDict(extra="ignore")
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="ignore")
 
     id: UUID
     name: str

prefect/client/schemas/schedules.py

@@ -3,13 +3,13 @@ Schedule schemas
 """
 
 import datetime
-from typing import Annotated, Any, Optional, Union
+from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Optional, Union
 
 import dateutil
 import dateutil.rrule
+import dateutil.tz
 import pendulum
 from pydantic import AfterValidator, ConfigDict, Field, field_validator, model_validator
-from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import TypeAlias, TypeGuard
 
 from prefect._internal.schemas.bases import PrefectBaseModel
@@ -20,6 +20,14 @@ from prefect._internal.schemas.validators import (
     validate_rrule_string,
 )
 
+if TYPE_CHECKING:
+    # type checkers have difficulty accepting that
+    # pydantic_extra_types.pendulum_dt and pendulum.DateTime can be used
+    # together.
+    DateTime = pendulum.DateTime
+else:
+    from pydantic_extra_types.pendulum_dt import DateTime
+
 MAX_ITERATIONS = 1000
 # approx. 1 years worth of RDATEs + buffer
 MAX_RRULE_LENGTH = 6500
@@ -54,7 +62,7 @@ class IntervalSchedule(PrefectBaseModel):
         timezone (str, optional): a valid timezone string
     """
 
-    model_config = ConfigDict(extra="forbid", exclude_none=True)
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
 
     interval: datetime.timedelta = Field(gt=datetime.timedelta(0))
     anchor_date: Annotated[DateTime, AfterValidator(default_anchor_date)] = Field(
@@ -68,6 +76,19 @@ class IntervalSchedule(PrefectBaseModel):
         self.timezone = default_timezone(self.timezone, self.model_dump())
         return self
 
+    if TYPE_CHECKING:
+        # The model accepts str or datetime values for `anchor_date`
+        def __init__(
+            self,
+            /,
+            interval: datetime.timedelta,
+            anchor_date: Optional[
+                Union[pendulum.DateTime, datetime.datetime, str]
+            ] = None,
+            timezone: Optional[str] = None,
+        ) -> None:
+            ...
+
 
 class CronSchedule(PrefectBaseModel):
     """
@@ -94,7 +115,7 @@ class CronSchedule(PrefectBaseModel):
 
     """
 
-    model_config = ConfigDict(extra="forbid")
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
 
     cron: str = Field(default=..., examples=["0 0 * * *"])
     timezone: Optional[str] = Field(default=None, examples=["America/New_York"])
@@ -107,18 +128,36 @@ class CronSchedule(PrefectBaseModel):
 
     @field_validator("timezone")
     @classmethod
-    def valid_timezone(cls, v):
+    def valid_timezone(cls, v: Optional[str]) -> str:
         return default_timezone(v)
 
     @field_validator("cron")
     @classmethod
-    def valid_cron_string(cls, v):
+    def valid_cron_string(cls, v: str) -> str:
         return validate_cron_string(v)
 
 
 DEFAULT_ANCHOR_DATE = pendulum.date(2020, 1, 1)
 
 
+def _rrule_dt(
+    rrule: dateutil.rrule.rrule, name: str = "_dtstart"
+) -> Optional[datetime.datetime]:
+    return getattr(rrule, name, None)
+
+
+def _rrule(
+    rruleset: dateutil.rrule.rruleset, name: str = "_rrule"
+) -> list[dateutil.rrule.rrule]:
+    return getattr(rruleset, name, [])
+
+
+def _rdates(
+    rrule: dateutil.rrule.rruleset, name: str = "_rdate"
+) -> list[datetime.datetime]:
+    return getattr(rrule, name, [])
+
+
 class RRuleSchedule(PrefectBaseModel):
     """
     RRule schedule, based on the iCalendar standard
@@ -139,7 +178,7 @@ class RRuleSchedule(PrefectBaseModel):
         timezone (str, optional): a valid timezone string
     """
 
-    model_config = ConfigDict(extra="forbid")
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
 
     rrule: str
     timezone: Optional[str] = Field(
@@ -148,58 +187,60 @@ class RRuleSchedule(PrefectBaseModel):
 
     @field_validator("rrule")
     @classmethod
-    def validate_rrule_str(cls, v):
+    def validate_rrule_str(cls, v: str) -> str:
         return validate_rrule_string(v)
 
     @classmethod
-    def from_rrule(cls, rrule: dateutil.rrule.rrule):
+    def from_rrule(
+        cls, rrule: Union[dateutil.rrule.rrule, dateutil.rrule.rruleset]
+    ) -> "RRuleSchedule":
         if isinstance(rrule, dateutil.rrule.rrule):
-            if rrule._dtstart.tzinfo is not None:
-                timezone = rrule._dtstart.tzinfo.name
+            dtstart = _rrule_dt(rrule)
+            if dtstart and dtstart.tzinfo is not None:
+                timezone = dtstart.tzinfo.tzname(dtstart)
             else:
                 timezone = "UTC"
             return RRuleSchedule(rrule=str(rrule), timezone=timezone)
-        elif isinstance(rrule, dateutil.rrule.rruleset):
-            dtstarts = [rr._dtstart for rr in rrule._rrule if rr._dtstart is not None]
-            unique_dstarts = set(pendulum.instance(d).in_tz("UTC") for d in dtstarts)
-            unique_timezones = set(d.tzinfo for d in dtstarts if d.tzinfo is not None)
-
-            if len(unique_timezones) > 1:
-                raise ValueError(
-                    f"rruleset has too many dtstart timezones: {unique_timezones}"
-                )
-
-            if len(unique_dstarts) > 1:
-                raise ValueError(f"rruleset has too many dtstarts: {unique_dstarts}")
-
-            if unique_dstarts and unique_timezones:
-                timezone = dtstarts[0].tzinfo.name
-            else:
-                timezone = "UTC"
-
-            rruleset_string = ""
-            if rrule._rrule:
-                rruleset_string += "\n".join(str(r) for r in rrule._rrule)
-            if rrule._exrule:
-                rruleset_string += "\n" if rruleset_string else ""
-                rruleset_string += "\n".join(str(r) for r in rrule._exrule).replace(
-                    "RRULE", "EXRULE"
-                )
-            if rrule._rdate:
-                rruleset_string += "\n" if rruleset_string else ""
-                rruleset_string += "RDATE:" + ",".join(
-                    rd.strftime("%Y%m%dT%H%M%SZ") for rd in rrule._rdate
-                )
-            if rrule._exdate:
-                rruleset_string += "\n" if rruleset_string else ""
-                rruleset_string += "EXDATE:" + ",".join(
-                    exd.strftime("%Y%m%dT%H%M%SZ") for exd in rrule._exdate
-                )
-            return RRuleSchedule(rrule=rruleset_string, timezone=timezone)
+        rrules = _rrule(rrule)
+        dtstarts = [dts for rr in rrules if (dts := _rrule_dt(rr)) is not None]
+        unique_dstarts = set(pendulum.instance(d).in_tz("UTC") for d in dtstarts)
+        unique_timezones = set(d.tzinfo for d in dtstarts if d.tzinfo is not None)
+
+        if len(unique_timezones) > 1:
+            raise ValueError(
+                f"rruleset has too many dtstart timezones: {unique_timezones}"
+            )
+
+        if len(unique_dstarts) > 1:
+            raise ValueError(f"rruleset has too many dtstarts: {unique_dstarts}")
+
+        if unique_dstarts and unique_timezones:
+            [unique_tz] = unique_timezones
+            timezone = unique_tz.tzname(dtstarts[0])
         else:
-            raise ValueError(f"Invalid RRule object: {rrule}")
-
-    def to_rrule(self) -> dateutil.rrule.rrule:
+            timezone = "UTC"
+
+        rruleset_string = ""
+        if rrules:
+            rruleset_string += "\n".join(str(r) for r in rrules)
+        if exrule := _rrule(rrule, "_exrule"):
+            rruleset_string += "\n" if rruleset_string else ""
+            rruleset_string += "\n".join(str(r) for r in exrule).replace(
+                "RRULE", "EXRULE"
+            )
+        if rdates := _rdates(rrule):
+            rruleset_string += "\n" if rruleset_string else ""
+            rruleset_string += "RDATE:" + ",".join(
+                rd.strftime("%Y%m%dT%H%M%SZ") for rd in rdates
+            )
+        if exdates := _rdates(rrule, "_exdate"):
+            rruleset_string += "\n" if rruleset_string else ""
+            rruleset_string += "EXDATE:" + ",".join(
+                exd.strftime("%Y%m%dT%H%M%SZ") for exd in exdates
+            )
+        return RRuleSchedule(rrule=rruleset_string, timezone=timezone)
+
+    def to_rrule(self) -> Union[dateutil.rrule.rrule, dateutil.rrule.rruleset]:
         """
         Since rrule doesn't properly serialize/deserialize timezones, we localize dates
         here
@@ -211,51 +252,53 @@ class RRuleSchedule(PrefectBaseModel):
         )
         timezone = dateutil.tz.gettz(self.timezone)
         if isinstance(rrule, dateutil.rrule.rrule):
-            kwargs = dict(dtstart=rrule._dtstart.replace(tzinfo=timezone))
-            if rrule._until:
+            dtstart = _rrule_dt(rrule)
+            assert dtstart is not None
+            kwargs: dict[str, Any] = dict(dtstart=dtstart.replace(tzinfo=timezone))
+            if until := _rrule_dt(rrule, "_until"):
                 kwargs.update(
-                    until=rrule._until.replace(tzinfo=timezone),
+                    until=until.replace(tzinfo=timezone),
                 )
             return rrule.replace(**kwargs)
-        elif isinstance(rrule, dateutil.rrule.rruleset):
-            # update rrules
-            localized_rrules = []
-            for rr in rrule._rrule:
-                kwargs = dict(dtstart=rr._dtstart.replace(tzinfo=timezone))
-                if rr._until:
-                    kwargs.update(
-                        until=rr._until.replace(tzinfo=timezone),
-                    )
-                localized_rrules.append(rr.replace(**kwargs))
-            rrule._rrule = localized_rrules
-
-            # update exrules
-            localized_exrules = []
-            for exr in rrule._exrule:
-                kwargs = dict(dtstart=exr._dtstart.replace(tzinfo=timezone))
-                if exr._until:
-                    kwargs.update(
-                        until=exr._until.replace(tzinfo=timezone),
-                    )
-                localized_exrules.append(exr.replace(**kwargs))
-            rrule._exrule = localized_exrules
-
-            # update rdates
-            localized_rdates = []
-            for rd in rrule._rdate:
-                localized_rdates.append(rd.replace(tzinfo=timezone))
-            rrule._rdate = localized_rdates
-
-            # update exdates
-            localized_exdates = []
-            for exd in rrule._exdate:
-                localized_exdates.append(exd.replace(tzinfo=timezone))
-            rrule._exdate = localized_exdates
-
-            return rrule
+
+        # update rrules
+        localized_rrules: list[dateutil.rrule.rrule] = []
+        for rr in _rrule(rrule):
+            dtstart = _rrule_dt(rr)
+            assert dtstart is not None
+            kwargs: dict[str, Any] = dict(dtstart=dtstart.replace(tzinfo=timezone))
+            if until := _rrule_dt(rr, "_until"):
+                kwargs.update(until=until.replace(tzinfo=timezone))
+            localized_rrules.append(rr.replace(**kwargs))
+        setattr(rrule, "_rrule", localized_rrules)
+
+        # update exrules
+        localized_exrules: list[dateutil.rrule.rruleset] = []
+        for exr in _rrule(rrule, "_exrule"):
+            dtstart = _rrule_dt(exr)
+            assert dtstart is not None
+            kwargs = dict(dtstart=dtstart.replace(tzinfo=timezone))
+            if until := _rrule_dt(exr, "_until"):
+                kwargs.update(until=until.replace(tzinfo=timezone))
+            localized_exrules.append(exr.replace(**kwargs))
+        setattr(rrule, "_exrule", localized_exrules)
+
+        # update rdates
+        localized_rdates: list[datetime.datetime] = []
+        for rd in _rdates(rrule):
+            localized_rdates.append(rd.replace(tzinfo=timezone))
+        setattr(rrule, "_rdate", localized_rdates)
+
+        # update exdates
+        localized_exdates: list[datetime.datetime] = []
+        for exd in _rdates(rrule, "_exdate"):
+            localized_exdates.append(exd.replace(tzinfo=timezone))
+        setattr(rrule, "_exdate", localized_exdates)
+
+        return rrule
 
     @field_validator("timezone")
-    def valid_timezone(cls, v):
+    def valid_timezone(cls, v: Optional[str]) -> str:
         """
         Validate that the provided timezone is a valid IANA timezone.
 
@@ -277,7 +320,7 @@ class RRuleSchedule(PrefectBaseModel):
 
 
 class NoSchedule(PrefectBaseModel):
-    model_config = ConfigDict(extra="forbid")
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
 
 
 SCHEDULE_TYPES: TypeAlias = Union[
@@ -326,7 +369,7 @@ def construct_schedule(
         if isinstance(interval, (int, float)):
             interval = datetime.timedelta(seconds=interval)
         if not anchor_date:
-            anchor_date = DateTime.now()
+            anchor_date = pendulum.DateTime.now()
         schedule = IntervalSchedule(
             interval=interval, anchor_date=anchor_date, timezone=timezone
         )

prefect/client/subscriptions.py

@@ -1,5 +1,7 @@
 import asyncio
-from typing import Any, Dict, Generic, Iterable, Optional, Type, TypeVar
+from collections.abc import Iterable
+from logging import Logger
+from typing import Any, Generic, Optional, TypeVar
 
 import orjson
 import websockets
@@ -11,7 +13,7 @@ from prefect._internal.schemas.bases import IDBaseModel
 from prefect.logging import get_logger
 from prefect.settings import PREFECT_API_KEY
 
-logger = get_logger(__name__)
+logger: Logger = get_logger(__name__)
 
 S = TypeVar("S", bound=IDBaseModel)
 
@@ -19,7 +21,7 @@ S = TypeVar("S", bound=IDBaseModel)
 class Subscription(Generic[S]):
     def __init__(
         self,
-        model: Type[S],
+        model: type[S],
         path: str,
         keys: Iterable[str],
         client_id: Optional[str] = None,
@@ -27,27 +29,33 @@ class Subscription(Generic[S]):
     ):
         self.model = model
         self.client_id = client_id
-        base_url = base_url.replace("http", "ws", 1)
-        self.subscription_url = f"{base_url}{path}"
+        base_url = base_url.replace("http", "ws", 1) if base_url else None
+        self.subscription_url: str = f"{base_url}{path}"
 
-        self.keys = list(keys)
+        self.keys: list[str] = list(keys)
 
         self._connect = websockets.connect(
             self.subscription_url,
-            subprotocols=["prefect"],
+            subprotocols=[websockets.Subprotocol("prefect")],
         )
         self._websocket = None
 
     def __aiter__(self) -> Self:
         return self
 
+    @property
+    def websocket(self) -> websockets.WebSocketClientProtocol:
+        if not self._websocket:
+            raise RuntimeError("Subscription is not connected")
+        return self._websocket
+
     async def __anext__(self) -> S:
         while True:
             try:
                 await self._ensure_connected()
-                message = await self._websocket.recv()
+                message = await self.websocket.recv()
 
-                await self._websocket.send(orjson.dumps({"type": "ack"}).decode())
+                await self.websocket.send(orjson.dumps({"type": "ack"}).decode())
 
                 return self.model.model_validate_json(message)
             except (
@@ -72,10 +80,10 @@ class Subscription(Generic[S]):
                 ).decode()
             )
 
-            auth: Dict[str, Any] = orjson.loads(await websocket.recv())
+            auth: dict[str, Any] = orjson.loads(await websocket.recv())
             assert auth["type"] == "auth_success", auth.get("message")
 
-            message = {"type": "subscribe", "keys": self.keys}
+            message: dict[str, Any] = {"type": "subscribe", "keys": self.keys}
             if self.client_id:
                 message.update({"client_id": self.client_id})
 
@@ -84,13 +92,19 @@ class Subscription(Generic[S]):
             AssertionError,
             websockets.exceptions.ConnectionClosedError,
         ) as e:
-            if isinstance(e, AssertionError) or e.rcvd.code == WS_1008_POLICY_VIOLATION:
+            if isinstance(e, AssertionError) or (
+                e.rcvd and e.rcvd.code == WS_1008_POLICY_VIOLATION
+            ):
                 if isinstance(e, AssertionError):
                     reason = e.args[0]
-                elif isinstance(e, websockets.exceptions.ConnectionClosedError):
+                elif e.rcvd and e.rcvd.reason:
                     reason = e.rcvd.reason
+                else:
+                    reason = "unknown"
+            else:
+                reason = None
 
-            if isinstance(e, AssertionError) or e.rcvd.code == WS_1008_POLICY_VIOLATION:
+            if reason:
                 raise Exception(
                     "Unable to authenticate to the subscription. Please "
                     "ensure the provided `PREFECT_API_KEY` you are using is "