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

prefect/types/_datetime.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import pendulum
+from pendulum.date import Date as PendulumDate
+from pendulum.datetime import DateTime as PendulumDateTime
+from pendulum.duration import Duration as PendulumDuration
+from pendulum.time import Time as PendulumTime
+from pydantic_extra_types.pendulum_dt import Date as PydanticDate
+from pydantic_extra_types.pendulum_dt import DateTime as PydanticDateTime
+from typing_extensions import TypeAlias
+
+DateTime: TypeAlias = PydanticDateTime
+Date: TypeAlias = PydanticDate
+
+
+def parse_datetime(
+    value: str,
+) -> PendulumDateTime | PendulumDate | PendulumTime | PendulumDuration:
+    return pendulum.parse(value)

prefect/utilities/_deprecated.py

@@ -0,0 +1,38 @@
+from typing import Any
+
+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 GetJsonSchemaHandler
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema
+
+
+class JsonPatch(JsonPatchBase):
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetJsonSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.typed_dict_schema(
+            {"patch": core_schema.typed_dict_field(core_schema.dict_schema())}
+        )
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> JsonSchemaValue:
+        json_schema = handler(core_schema)
+        json_schema = handler.resolve_ref_schema(json_schema)
+        json_schema.pop("required", None)
+        json_schema.pop("properties", None)
+        json_schema.update(
+            {
+                "type": "array",
+                "format": "rfc6902",
+                "items": {
+                    "type": "object",
+                    "additionalProperties": {"type": "string"},
+                },
+            }
+        )
+        return json_schema

prefect/utilities/engine.py

@@ -757,12 +757,19 @@ def resolve_to_final_result(expr: Any, context: dict[str, Any]) -> Any:
         parameter_context = propagate.extract(
             {"traceparent": state.state_details.traceparent}
         )
-        trace.get_current_span().add_link(
-            context=trace.get_current_span(parameter_context).get_span_context(),
-            attributes={
+        attributes = {}
+
+        # If this future is being used as a parameter (as opposed to just a wait_for),
+        # add attributes to the span to indicate the parameter name and type
+        if "parameter_name" in context:
+            attributes = {
                 "prefect.input.name": context["parameter_name"],
                 "prefect.input.type": type(result).__name__,
-            },
+            }
+
+        trace.get_current_span().add_link(
+            context=trace.get_current_span(parameter_context).get_span_context(),
+            attributes=attributes,
         )
 
     return result

prefect/utilities/filesystem.py

@@ -8,7 +8,7 @@ import threading
 from collections.abc import Iterable
 from contextlib import contextmanager
 from pathlib import Path, PureWindowsPath
-from typing import AnyStr, Optional, Union, cast
+from typing import Any, AnyStr, Optional, Union, cast
 
 # fsspec has no stubs, see https://github.com/fsspec/filesystem_spec/issues/625
 import fsspec  # type: ignore
@@ -114,7 +114,7 @@ def filename(path: str) -> str:
     return path.split(sep)[-1]
 
 
-def is_local_path(path: Union[str, pathlib.Path, OpenFile]) -> bool:
+def is_local_path(path: Union[str, pathlib.Path, Any]) -> bool:
     """Check if the given path points to a local or remote file system"""
     if isinstance(path, str):
         try:

prefect/utilities/generics.py

@@ -5,7 +5,7 @@ from pydantic_core import SchemaValidator, core_schema
 
 T = TypeVar("T", bound=BaseModel)
 
-ListValidator = SchemaValidator(
+ListValidator: SchemaValidator = SchemaValidator(
     schema=core_schema.list_schema(
         items_schema=core_schema.dict_schema(
             keys_schema=core_schema.str_schema(), values_schema=core_schema.any_schema()

prefect/utilities/pydantic.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import (
     Any,
     Callable,
@@ -10,18 +11,13 @@ from typing import (
     overload,
 )
 
-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,
     Secret,
     TypeAdapter,
     ValidationError,
 )
-from pydantic.json_schema import JsonSchemaValue
-from pydantic_core import core_schema, to_jsonable_python
+from pydantic_core import to_jsonable_python
 from typing_extensions import Literal
 
 from prefect.utilities.dispatch import get_dispatch_key, lookup_type, register_base_type
@@ -262,36 +258,6 @@ class PartialModel(Generic[M]):
         return f"PartialModel(cls={self.model_cls.__name__}, {dsp_fields})"
 
 
-class JsonPatch(JsonPatchBase):
-    @classmethod
-    def __get_pydantic_core_schema__(
-        cls, source_type: Any, handler: GetJsonSchemaHandler
-    ) -> core_schema.CoreSchema:
-        return core_schema.typed_dict_schema(
-            {"patch": core_schema.typed_dict_field(core_schema.dict_schema())}
-        )
-
-    @classmethod
-    def __get_pydantic_json_schema__(
-        cls, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
-    ) -> JsonSchemaValue:
-        json_schema = handler(core_schema)
-        json_schema = handler.resolve_ref_schema(json_schema)
-        json_schema.pop("required", None)
-        json_schema.pop("properties", None)
-        json_schema.update(
-            {
-                "type": "array",
-                "format": "rfc6902",
-                "items": {
-                    "type": "object",
-                    "additionalProperties": {"type": "string"},
-                },
-            }
-        )
-        return json_schema
-
-
 def custom_pydantic_encoder(
     type_encoders: dict[Any, Callable[[type[Any]], Any]], obj: Any
 ) -> Any:
@@ -382,3 +348,22 @@ def handle_secret_render(value: object, context: dict[str, Any]) -> object:
     elif isinstance(value, BaseModel):
         return value.model_dump(context=context)
     return value
+
+
+def __getattr__(name: str) -> Any:
+    """
+    Handles imports from this module that are deprecated.
+    """
+
+    if name == "JsonPatch":
+        warnings.warn(
+            "JsonPatch is deprecated and will be removed after March 2025. "
+            "Please use `JsonPatch` from the `jsonpatch` package instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        from ._deprecated import JsonPatch
+
+        return JsonPatch
+    else:
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

prefect/workers/base.py

@@ -1,9 +1,22 @@
+from __future__ import annotations
+
 import abc
 import asyncio
 import threading
 from contextlib import AsyncExitStack
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Set,
+    Type,
+    Union,
+)
 from uuid import UUID, uuid4
 
 import anyio
@@ -13,7 +26,7 @@ import pendulum
 from importlib_metadata import distributions
 from pydantic import BaseModel, Field, PrivateAttr, field_validator
 from pydantic.json_schema import GenerateJsonSchema
-from typing_extensions import Literal
+from typing_extensions import Literal, Self, TypeVar
 
 import prefect
 from prefect._internal.schemas.validators import return_v_or_none
@@ -104,7 +117,7 @@ class BaseJobConfiguration(BaseModel):
     _related_objects: Dict[str, Any] = PrivateAttr(default_factory=dict)
 
     @property
-    def is_using_a_runner(self):
+    def is_using_a_runner(self) -> bool:
         return self.command is not None and "prefect flow-run execute" in self.command
 
     @field_validator("command")
@@ -175,7 +188,7 @@ class BaseJobConfiguration(BaseModel):
         return cls(**populated_configuration)
 
     @classmethod
-    def json_template(cls) -> dict:
+    def json_template(cls) -> dict[str, Any]:
         """Returns a dict with job configuration as keys and the corresponding templates as values
 
         Defaults to using the job configuration parameter name as the template variable name.
@@ -186,7 +199,7 @@ class BaseJobConfiguration(BaseModel):
             key2: '{{ template2 }}', # `template2` specifically provide as template
         }
         """
-        configuration = {}
+        configuration: dict[str, Any] = {}
         properties = cls.model_json_schema()["properties"]
         for k, v in properties.items():
             if v.get("template"):
@@ -202,7 +215,7 @@ class BaseJobConfiguration(BaseModel):
         flow_run: "FlowRun",
         deployment: Optional["DeploymentResponse"] = None,
         flow: Optional["Flow"] = None,
-    ):
+    ) -> None:
         """
         Prepare the job configuration for a flow run.
 
@@ -368,15 +381,20 @@ class BaseWorkerResult(BaseModel, abc.ABC):
     identifier: str
     status_code: int
 
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return self.status_code == 0
 
 
+C = TypeVar("C", bound=BaseJobConfiguration)
+V = TypeVar("V", bound=BaseVariables)
+R = TypeVar("R", bound=BaseWorkerResult)
+
+
 @register_base_type
-class BaseWorker(abc.ABC):
+class BaseWorker(abc.ABC, Generic[C, V, R]):
     type: str
-    job_configuration: Type[BaseJobConfiguration] = BaseJobConfiguration
-    job_configuration_variables: Optional[Type[BaseVariables]] = None
+    job_configuration: Type[C] = BaseJobConfiguration  # type: ignore
+    job_configuration_variables: Optional[Type[V]] = None
 
     _documentation_url = ""
     _logo_url = ""
@@ -418,7 +436,7 @@ class BaseWorker(abc.ABC):
         """
         if name and ("/" in name or "%" in name):
             raise ValueError("Worker name cannot contain '/' or '%'")
-        self.name = name or f"{self.__class__.__name__} {uuid4()}"
+        self.name: str = name or f"{self.__class__.__name__} {uuid4()}"
         self._started_event: Optional[Event] = None
         self.backend_id: Optional[UUID] = None
         self._logger = get_worker_logger(self)
@@ -432,7 +450,7 @@ class BaseWorker(abc.ABC):
         self._prefetch_seconds: float = (
             prefetch_seconds or PREFECT_WORKER_PREFETCH_SECONDS.value()
         )
-        self.heartbeat_interval_seconds = (
+        self.heartbeat_interval_seconds: int = (
             heartbeat_interval_seconds or PREFECT_WORKER_HEARTBEAT_SECONDS.value()
         )
 
@@ -461,7 +479,7 @@ class BaseWorker(abc.ABC):
         return cls._description
 
     @classmethod
-    def get_default_base_job_template(cls) -> Dict:
+    def get_default_base_job_template(cls) -> dict[str, Any]:
         if cls.job_configuration_variables is None:
             schema = cls.job_configuration.model_json_schema()
             # remove "template" key from all dicts in schema['properties'] because it is not a
@@ -479,7 +497,9 @@ class BaseWorker(abc.ABC):
         }
 
     @staticmethod
-    def get_worker_class_from_type(type: str) -> Optional[Type["BaseWorker"]]:
+    def get_worker_class_from_type(
+        type: str,
+    ) -> Optional[Type["BaseWorker[Any, Any, Any]"]]:
         """
         Returns the worker class for a given worker type. If the worker type
         is not recognized, returns None.
@@ -500,7 +520,7 @@ class BaseWorker(abc.ABC):
             return list(worker_registry.keys())
         return []
 
-    def get_name_slug(self):
+    def get_name_slug(self) -> str:
         return slugify(self.name)
 
     def get_flow_run_logger(self, flow_run: "FlowRun") -> PrefectLogAdapter:
@@ -524,7 +544,7 @@ class BaseWorker(abc.ABC):
         run_once: bool = False,
         with_healthcheck: bool = False,
         printer: Callable[..., None] = print,
-    ):
+    ) -> None:
         """
         Starts the worker and runs the main worker loops.
 
@@ -603,9 +623,9 @@ class BaseWorker(abc.ABC):
     async def run(
         self,
         flow_run: "FlowRun",
-        configuration: BaseJobConfiguration,
-        task_status: Optional[anyio.abc.TaskStatus] = None,
-    ) -> BaseWorkerResult:
+        configuration: C,
+        task_status: Optional[anyio.abc.TaskStatus[int]] = None,
+    ) -> R:
         """
         Runs a given flow run on the current worker.
         """
@@ -614,12 +634,12 @@ class BaseWorker(abc.ABC):
         )
 
     @classmethod
-    def __dispatch_key__(cls):
+    def __dispatch_key__(cls) -> str | None:
         if cls.__name__ == "BaseWorker":
             return None  # The base class is abstract
         return cls.type
 
-    async def setup(self):
+    async def setup(self) -> None:
         """Prepares the worker to run."""
         self._logger.debug("Setting up worker...")
         self._runs_task_group = anyio.create_task_group()
@@ -637,10 +657,10 @@ class BaseWorker(abc.ABC):
 
         self.is_setup = True
 
-    async def teardown(self, *exc_info):
+    async def teardown(self, *exc_info: Any) -> None:
         """Cleans up resources after the worker is stopped."""
         self._logger.debug("Tearing down worker...")
-        self.is_setup = False
+        self.is_setup: bool = False
         for scope in self._scheduled_task_scopes:
             scope.cancel()
 
@@ -684,14 +704,16 @@ class BaseWorker(abc.ABC):
 
         return is_still_polling
 
-    async def get_and_submit_flow_runs(self):
+    async def get_and_submit_flow_runs(self) -> list["FlowRun"]:
         runs_response = await self._get_scheduled_flow_runs()
 
         self._last_polled_time = pendulum.now("utc")
 
         return await self._submit_scheduled_flow_runs(flow_run_response=runs_response)
 
-    async def _update_local_work_pool_info(self):
+    async def _update_local_work_pool_info(self) -> None:
+        if TYPE_CHECKING:
+            assert self._client is not None
         try:
             work_pool = await self._client.read_work_pool(
                 work_pool_name=self._work_pool_name
@@ -803,7 +825,7 @@ class BaseWorker(abc.ABC):
 
         return worker_id
 
-    async def sync_with_backend(self):
+    async def sync_with_backend(self) -> None:
         """
         Updates the worker's local information about it's current work pool and
         queues. Sends a worker heartbeat to the API.
@@ -1042,7 +1064,7 @@ class BaseWorker(abc.ABC):
             self._limiter.release_on_behalf_of(flow_run_id)
             self._logger.debug("Limit slot released for flow run '%s'", flow_run_id)
 
-    def get_status(self):
+    def get_status(self) -> dict[str, Any]:
         """
         Retrieves the status of the current worker including its name, current worker
         pool, the work pool queues it is polling, and its local settings.
@@ -1234,17 +1256,17 @@ class BaseWorker(abc.ABC):
 
             await self._client.update_flow_run_labels(flow_run_id, labels)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Self:
         self._logger.debug("Entering worker context...")
         await self.setup()
 
         return self
 
-    async def __aexit__(self, *exc_info):
+    async def __aexit__(self, *exc_info: Any) -> None:
         self._logger.debug("Exiting worker context...")
         await self.teardown(*exc_info)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"Worker(pool={self._work_pool_name!r}, name={self.name!r})"
 
     def _event_resource(self):

prefect/workers/process.py

@@ -13,6 +13,7 @@ to poll for flow runs.
 For more information about work pools and workers,
 checkout out the [Prefect docs](/concepts/work-pools/).
 """
+from __future__ import annotations
 
 import contextlib
 import os
@@ -30,7 +31,7 @@ import anyio
 import anyio.abc
 from pydantic import Field, field_validator
 
-from prefect._internal.schemas.validators import validate_command
+from prefect._internal.schemas.validators import validate_working_dir
 from prefect.client.schemas import FlowRun
 from prefect.client.schemas.filters import (
     FlowRunFilter,
@@ -85,19 +86,21 @@ class ProcessJobConfiguration(BaseJobConfiguration):
 
     @field_validator("working_dir")
     @classmethod
-    def validate_command(cls, v: str) -> str:
-        return validate_command(v)
+    def validate_working_dir(cls, v: Path | str | None) -> Path | None:
+        if isinstance(v, str):
+            return validate_working_dir(v)
+        return v
 
     def prepare_for_flow_run(
         self,
         flow_run: "FlowRun",
         deployment: Optional["DeploymentResponse"] = None,
         flow: Optional["Flow"] = None,
-    ):
+    ) -> None:
         super().prepare_for_flow_run(flow_run, deployment, flow)
 
-        self.env = {**os.environ, **self.env}
-        self.command = (
+        self.env: dict[str, str | None] = {**os.environ, **self.env}
+        self.command: str | None = (
             f"{get_sys_executable()} -m prefect.engine"
             if self.command == self._base_flow_run_command()
             else self.command
@@ -134,10 +137,12 @@ class ProcessWorkerResult(BaseWorkerResult):
     """Contains information about the final state of a completed process"""
 
 
-class ProcessWorker(BaseWorker):
+class ProcessWorker(
+    BaseWorker[ProcessJobConfiguration, ProcessVariables, ProcessWorkerResult]
+):
     type = "process"
-    job_configuration = ProcessJobConfiguration
-    job_configuration_variables = ProcessVariables
+    job_configuration: type[ProcessJobConfiguration] = ProcessJobConfiguration
+    job_configuration_variables: type[ProcessVariables] | None = ProcessVariables
 
     _description = (
         "Execute flow runs as subprocesses on a worker. Works well for local execution"
@@ -152,7 +157,7 @@ class ProcessWorker(BaseWorker):
         run_once: bool = False,
         with_healthcheck: bool = False,
         printer: Callable[..., None] = print,
-    ):
+    ) -> None:
         """
         Starts the worker and runs the main worker loops.
 
@@ -241,8 +246,8 @@ class ProcessWorker(BaseWorker):
         self,
         flow_run: FlowRun,
         configuration: ProcessJobConfiguration,
-        task_status: Optional[anyio.abc.TaskStatus] = None,
-    ):
+        task_status: Optional[anyio.abc.TaskStatus[int]] = None,
+    ) -> ProcessWorkerResult:
         command = configuration.command
         if not command:
             command = f"{get_sys_executable()} -m prefect.engine"
@@ -322,7 +327,7 @@ class ProcessWorker(BaseWorker):
         self,
         infrastructure_pid: str,
         grace_seconds: int = 30,
-    ):
+    ) -> None:
         hostname, pid = _parse_infrastructure_pid(infrastructure_pid)
 
         if hostname != socket.gethostname():
@@ -372,7 +377,7 @@ class ProcessWorker(BaseWorker):
                 # process ended right after the check above.
                 return
 
-    async def check_for_cancelled_flow_runs(self):
+    async def check_for_cancelled_flow_runs(self) -> list["FlowRun"]:
         if not self.is_setup:
             raise RuntimeError(
                 "Worker is not set up. Please make sure you are running this worker "
@@ -429,7 +434,7 @@ class ProcessWorker(BaseWorker):
 
         return cancelling_flow_runs
 
-    async def cancel_run(self, flow_run: "FlowRun"):
+    async def cancel_run(self, flow_run: "FlowRun") -> None:
         run_logger = self.get_flow_run_logger(flow_run)
 
         try:

prefect/workers/server.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Any
 
 import uvicorn
 import uvicorn.server
@@ -10,14 +10,13 @@ from prefect.settings import (
     PREFECT_WORKER_WEBSERVER_PORT,
 )
 from prefect.workers.base import BaseWorker
-from prefect.workers.process import ProcessWorker
 
 
 def build_healthcheck_server(
-    worker: Union[BaseWorker, ProcessWorker],
+    worker: BaseWorker[Any, Any, Any],
     query_interval_seconds: float,
     log_level: str = "error",
-):
+) -> uvicorn.Server:
     """
     Build a healthcheck FastAPI server for a worker.
 
@@ -54,7 +53,7 @@ def build_healthcheck_server(
 
 
 def start_healthcheck_server(
-    worker: Union[BaseWorker, ProcessWorker],
+    worker: BaseWorker[Any, Any, Any],
     query_interval_seconds: float,
     log_level: str = "error",
 ) -> None:

{prefect_client-3.1.12.dist-info → prefect_client-3.1.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.1.12
+Version: 3.1.14
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.
@@ -45,7 +45,7 @@ Requires-Dist: packaging<24.3,>=21.3
 Requires-Dist: pathspec>=0.8.0
 Requires-Dist: pendulum<4,>=3.0.0
 Requires-Dist: prometheus-client>=0.20.0
-Requires-Dist: pydantic!=2.10.0,<3.0.0,>=2.7
+Requires-Dist: pydantic!=2.10.0,<3.0.0,>=2.9
 Requires-Dist: pydantic-core<3.0.0,>=2.12.0
 Requires-Dist: pydantic-extra-types<3.0.0,>=2.8.2
 Requires-Dist: pydantic-settings>2.2.1