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

prefect/settings/models/testing.py

@@ -1,12 +1,13 @@
 from typing import Any, ClassVar, Optional
 
-from pydantic import AliasChoices, AliasPath, ConfigDict, Field
+from pydantic import AliasChoices, AliasPath, Field
+from pydantic_settings import SettingsConfigDict
 
-from prefect.settings.base import PrefectBaseSettings, _build_settings_config
+from prefect.settings.base import PrefectBaseSettings, build_settings_config
 
 
 class TestingSettings(PrefectBaseSettings):
-    model_config: ClassVar[ConfigDict] = _build_settings_config(("testing",))
+    model_config: ClassVar[SettingsConfigDict] = build_settings_config(("testing",))
 
     test_mode: bool = Field(
         default=False,

prefect/settings/models/worker.py

@@ -1,12 +1,15 @@
 from typing import ClassVar
 
-from pydantic import ConfigDict, Field
+from pydantic import Field
+from pydantic_settings import SettingsConfigDict
 
-from prefect.settings.base import PrefectBaseSettings, _build_settings_config
+from prefect.settings.base import PrefectBaseSettings, build_settings_config
 
 
 class WorkerWebserverSettings(PrefectBaseSettings):
-    model_config: ClassVar[ConfigDict] = _build_settings_config(("worker", "webserver"))
+    model_config: ClassVar[SettingsConfigDict] = build_settings_config(
+        ("worker", "webserver")
+    )
 
     host: str = Field(
         default="0.0.0.0",
@@ -20,7 +23,7 @@ class WorkerWebserverSettings(PrefectBaseSettings):
 
 
 class WorkerSettings(PrefectBaseSettings):
-    model_config: ClassVar[ConfigDict] = _build_settings_config(("worker",))
+    model_config: ClassVar[SettingsConfigDict] = build_settings_config(("worker",))
 
     heartbeat_seconds: float = Field(
         default=30,

prefect/settings/profiles.py

@@ -7,10 +7,9 @@ from typing import (
     ClassVar,
     Dict,
     Iterable,
-    List,
+    Iterator,
     Optional,
     Set,
-    Tuple,
     Union,
 )
 
@@ -77,8 +76,8 @@ class Profile(BaseModel):
             if value is not None
         }
 
-    def validate_settings(self):
-        errors: List[Tuple[Setting, ValidationError]] = []
+    def validate_settings(self) -> None:
+        errors: list[tuple[Setting, ValidationError]] = []
         for setting, value in self.settings.items():
             try:
                 model_fields = Settings.model_fields
@@ -109,7 +108,9 @@ class ProfilesCollection:
     def __init__(
         self, profiles: Iterable[Profile], active: Optional[str] = None
     ) -> None:
-        self.profiles_by_name = {profile.name: profile for profile in profiles}
+        self.profiles_by_name: dict[str, Profile] = {
+            profile.name: profile for profile in profiles
+        }
         self.active_name = active
 
     @property
@@ -128,7 +129,7 @@ class ProfilesCollection:
             return None
         return self[self.active_name]
 
-    def set_active(self, name: Optional[str], check: bool = True):
+    def set_active(self, name: Optional[str], check: bool = True) -> None:
         """
         Set the active profile name in the collection.
 
@@ -142,7 +143,7 @@ class ProfilesCollection:
     def update_profile(
         self,
         name: str,
-        settings: Dict[Setting, Any],
+        settings: dict[Setting, Any],
         source: Optional[Path] = None,
     ) -> Profile:
         """
@@ -214,7 +215,7 @@ class ProfilesCollection:
             active=self.active_name,
         )
 
-    def to_dict(self):
+    def to_dict(self) -> dict[str, Any]:
         """
         Convert to a dictionary suitable for writing to disk.
         """
@@ -229,11 +230,11 @@ class ProfilesCollection:
     def __getitem__(self, name: str) -> Profile:
         return self.profiles_by_name[name]
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[str]:
         return self.profiles_by_name.__iter__()
 
-    def items(self):
-        return self.profiles_by_name.items()
+    def items(self) -> list[tuple[str, Profile]]:
+        return list(self.profiles_by_name.items())
 
     def __eq__(self, __o: object) -> bool:
         if not isinstance(__o, ProfilesCollection):
@@ -325,7 +326,7 @@ def load_profiles(include_defaults: bool = True) -> ProfilesCollection:
     return profiles
 
 
-def load_current_profile():
+def load_current_profile() -> Profile:
     """
     Load the current profile from the default and current profile paths.
 

prefect/settings/sources.py

@@ -2,7 +2,7 @@ import os
 import sys
 import warnings
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Type
 
 import dotenv
 import toml
@@ -54,6 +54,7 @@ class EnvFilterSettingsSource(EnvSettingsSource):
             env_parse_none_str,
             env_parse_enums,
         )
+        self.env_vars: Mapping[str, str | None]
         if env_filter:
             if isinstance(self.env_vars, dict):
                 for key in env_filter:
@@ -97,7 +98,7 @@ class FilteredDotEnvSettingsSource(DotEnvSettingsSource):
                 for key in self.env_blacklist:
                     self.env_vars.pop(key, None)
             else:
-                self.env_vars = {
+                self.env_vars: dict[str, str | None] = {
                     key: value
                     for key, value in self.env_vars.items()  # type: ignore
                     if key.lower() not in env_blacklist
@@ -114,8 +115,8 @@ class ProfileSettingsTomlLoader(PydanticBaseSettingsSource):
     def __init__(self, settings_cls: Type[BaseSettings]):
         super().__init__(settings_cls)
         self.settings_cls = settings_cls
-        self.profiles_path = _get_profiles_path()
-        self.profile_settings = self._load_profile_settings()
+        self.profiles_path: Path = _get_profiles_path()
+        self.profile_settings: dict[str, Any] = self._load_profile_settings()
 
     def _load_profile_settings(self) -> Dict[str, Any]:
         """Helper method to load the profile settings from the profiles.toml file"""
@@ -213,14 +214,14 @@ class TomlConfigSettingsSourceBase(PydanticBaseSettingsSource, ConfigFileSourceM
     def __init__(self, settings_cls: Type[BaseSettings]):
         super().__init__(settings_cls)
         self.settings_cls = settings_cls
-        self.toml_data = {}
+        self.toml_data: dict[str, Any] = {}
 
-    def _read_file(self, path: Path) -> Dict[str, Any]:
+    def _read_file(self, path: Path) -> dict[str, Any]:
         return toml.load(path)
 
     def get_field_value(
         self, field: FieldInfo, field_name: str
-    ) -> Tuple[Any, str, bool]:
+    ) -> tuple[Any, str, bool]:
         """Concrete implementation to get the field value from toml data"""
         value = self.toml_data.get(field_name)
         if isinstance(value, dict):
@@ -244,9 +245,9 @@ class TomlConfigSettingsSourceBase(PydanticBaseSettingsSource, ConfigFileSourceM
                         break
         return value, name, self.field_is_complex(field)
 
-    def __call__(self) -> Dict[str, Any]:
+    def __call__(self) -> dict[str, Any]:
         """Called by pydantic to get the settings from our custom source"""
-        toml_setings: Dict[str, Any] = {}
+        toml_setings: dict[str, Any] = {}
         for field_name, field in self.settings_cls.model_fields.items():
             value, key, is_complex = self.get_field_value(field, field_name)
             if value is not None:
@@ -265,15 +266,15 @@ class PrefectTomlConfigSettingsSource(TomlConfigSettingsSourceBase):
         settings_cls: Type[BaseSettings],
     ):
         super().__init__(settings_cls)
-        self.toml_file_path = settings_cls.model_config.get(
-            "toml_file", DEFAULT_PREFECT_TOML_PATH
-        )
-        self.toml_data = self._read_files(self.toml_file_path)
-        self.toml_table_header = settings_cls.model_config.get(
+        self.toml_file_path: Path | str | Sequence[
+            Path | str
+        ] | None = settings_cls.model_config.get("toml_file", DEFAULT_PREFECT_TOML_PATH)
+        self.toml_data: dict[str, Any] = self._read_files(self.toml_file_path)
+        self.toml_table_header: tuple[str, ...] = settings_cls.model_config.get(
             "prefect_toml_table_header", tuple()
         )
         for key in self.toml_table_header:
-            self.toml_data = self.toml_data.get(key, {})
+            self.toml_data: dict[str, Any] = self.toml_data.get(key, {})
 
 
 class PyprojectTomlConfigSettingsSource(TomlConfigSettingsSourceBase):
@@ -284,13 +285,13 @@ class PyprojectTomlConfigSettingsSource(TomlConfigSettingsSourceBase):
         settings_cls: Type[BaseSettings],
     ):
         super().__init__(settings_cls)
-        self.toml_file_path = Path("pyproject.toml")
-        self.toml_data = self._read_files(self.toml_file_path)
-        self.toml_table_header = settings_cls.model_config.get(
+        self.toml_file_path: Path = Path("pyproject.toml")
+        self.toml_data: dict[str, Any] = self._read_files(self.toml_file_path)
+        self.toml_table_header: tuple[str, ...] = settings_cls.model_config.get(
             "pyproject_toml_table_header", ("tool", "prefect")
         )
         for key in self.toml_table_header:
-            self.toml_data = self.toml_data.get(key, {})
+            self.toml_data: dict[str, Any] = self.toml_data.get(key, {})
 
 
 def _is_test_mode() -> bool:

prefect/states.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import datetime
 import sys
@@ -33,12 +35,14 @@ from prefect.utilities.asyncutils import in_async_main_thread, sync_compatible
 from prefect.utilities.collections import ensure_iterable
 
 if TYPE_CHECKING:
+    import logging
+
     from prefect.results import (
         R,
         ResultStore,
     )
 
-logger = get_logger("states")
+logger: "logging.Logger" = get_logger("states")
 
 
 @deprecated.deprecated_parameter(
@@ -245,8 +249,8 @@ async def exception_to_failed_state(
     exc: Optional[BaseException] = None,
     result_store: Optional["ResultStore"] = None,
     write_result: bool = False,
-    **kwargs,
-) -> State:
+    **kwargs: Any,
+) -> State[BaseException]:
     """
     Convenience function for creating `Failed` states from exceptions
     """
@@ -553,17 +557,17 @@ def is_state_iterable(obj: Any) -> TypeGuard[Iterable[State]]:
 
 
 class StateGroup:
-    def __init__(self, states: Iterable[State]) -> None:
-        self.states = states
-        self.type_counts = self._get_type_counts(states)
-        self.total_count = len(states)
-        self.cancelled_count = self.type_counts[StateType.CANCELLED]
-        self.final_count = sum(state.is_final() for state in states)
-        self.not_final_count = self.total_count - self.final_count
-        self.paused_count = self.type_counts[StateType.PAUSED]
+    def __init__(self, states: list[State]) -> None:
+        self.states: list[State] = states
+        self.type_counts: dict[StateType, int] = self._get_type_counts(states)
+        self.total_count: int = len(states)
+        self.cancelled_count: int = self.type_counts[StateType.CANCELLED]
+        self.final_count: int = sum(state.is_final() for state in states)
+        self.not_final_count: int = self.total_count - self.final_count
+        self.paused_count: int = self.type_counts[StateType.PAUSED]
 
     @property
-    def fail_count(self):
+    def fail_count(self) -> int:
         return self.type_counts[StateType.FAILED] + self.type_counts[StateType.CRASHED]
 
     def all_completed(self) -> bool:
@@ -741,7 +745,7 @@ def Suspended(
     pause_expiration_time: Optional[datetime.datetime] = None,
     pause_key: Optional[str] = None,
     **kwargs: Any,
-):
+) -> "State[R]":
     """Convenience function for creating `Suspended` states.
 
     Returns:

prefect/task_engine.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import inspect
 import logging
@@ -28,8 +30,9 @@ from uuid import UUID
 import anyio
 import pendulum
 from opentelemetry import trace
-from typing_extensions import ParamSpec
+from typing_extensions import ParamSpec, Self
 
+from prefect.cache_policies import CachePolicy
 from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client
 from prefect.client.schemas import TaskRun
 from prefect.client.schemas.objects import State, TaskRunInput
@@ -55,7 +58,7 @@ from prefect.exceptions import (
 from prefect.logging.loggers import get_logger, patch_print, task_run_logger
 from prefect.results import (
     ResultRecord,
-    _format_user_supplied_storage_key,
+    _format_user_supplied_storage_key,  # type: ignore[reportPrivateUsage]
     get_result_store,
     should_persist_result,
 )
@@ -115,20 +118,20 @@ class BaseTaskRunEngine(Generic[P, R]):
     # holds the return value from the user code
     _return_value: Union[R, Type[NotSet]] = NotSet
     # holds the exception raised by the user code, if any
-    _raised: Union[Exception, Type[NotSet]] = NotSet
+    _raised: Union[Exception, BaseException, Type[NotSet]] = NotSet
     _initial_run_context: Optional[TaskRunContext] = None
     _is_started: bool = False
     _task_name_set: bool = False
     _last_event: Optional[PrefectEvent] = None
     _telemetry: RunTelemetry = field(default_factory=RunTelemetry)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.parameters is None:
             self.parameters = {}
 
     @property
     def state(self) -> State:
-        if not self.task_run:
+        if not self.task_run or not self.task_run.state:
             raise ValueError("Task run is not set")
         return self.task_run.state
 
@@ -142,8 +145,8 @@ class BaseTaskRunEngine(Generic[P, R]):
         return False
 
     def compute_transaction_key(self) -> Optional[str]:
-        key = None
-        if self.task.cache_policy:
+        key: Optional[str] = None
+        if self.task.cache_policy and isinstance(self.task.cache_policy, CachePolicy):
             flow_run_context = FlowRunContext.get()
             task_run_context = TaskRunContext.get()
 
@@ -153,10 +156,12 @@ class BaseTaskRunEngine(Generic[P, R]):
                 parameters = None
 
             try:
+                if not task_run_context:
+                    raise ValueError("Task run context is not set")
                 key = self.task.cache_policy.compute_key(
                     task_ctx=task_run_context,
-                    inputs=self.parameters,
-                    flow_parameters=parameters,
+                    inputs=self.parameters or {},
+                    flow_parameters=parameters or {},
                 )
             except Exception:
                 self.logger.exception(
@@ -169,7 +174,7 @@ class BaseTaskRunEngine(Generic[P, R]):
 
     def _resolve_parameters(self):
         if not self.parameters:
-            return {}
+            return None
 
         resolved_parameters = {}
         for parameter, value in self.parameters.items():
@@ -227,10 +232,8 @@ class BaseTaskRunEngine(Generic[P, R]):
         if self.task_run and self.task_run.start_time and not self.task_run.end_time:
             self.task_run.end_time = state.timestamp
 
-            if self.task_run.state.is_running():
-                self.task_run.total_run_time += (
-                    state.timestamp - self.task_run.state.timestamp
-                )
+            if self.state.is_running():
+                self.task_run.total_run_time += state.timestamp - self.state.timestamp
 
     def is_running(self) -> bool:
         """Whether or not the engine is currently running a task."""
@@ -238,7 +241,7 @@ class BaseTaskRunEngine(Generic[P, R]):
             return False
         return task_run.state.is_running() or task_run.state.is_scheduled()
 
-    def log_finished_message(self):
+    def log_finished_message(self) -> None:
         if not self.task_run:
             return
 
@@ -294,6 +297,7 @@ class BaseTaskRunEngine(Generic[P, R]):
 
 @dataclass
 class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
+    task_run: Optional[TaskRun] = None
     _client: Optional[SyncPrefectClient] = None
 
     @property
@@ -336,7 +340,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
             return False
 
-    def call_hooks(self, state: Optional[State] = None):
+    def call_hooks(self, state: Optional[State] = None) -> None:
         if state is None:
             state = self.state
         task = self.task
@@ -371,7 +375,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 self.logger.info(f"Hook {hook_name!r} finished running successfully")
 
-    def begin_run(self):
+    def begin_run(self) -> None:
         try:
             self._resolve_parameters()
             self._set_custom_task_run_name()
@@ -390,6 +394,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
         new_state = Running()
 
+        assert self.task_run is not None, "Task run is not set"
         self.task_run.start_time = new_state.timestamp
 
         flow_run_context = FlowRunContext.get()
@@ -406,7 +411,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         # result reference that no longer exists
         if state.is_completed():
             try:
-                state.result(retry_result_failure=False, _sync=True)
+                state.result(retry_result_failure=False, _sync=True)  # type: ignore[reportCallIssue]
             except Exception:
                 state = self.set_state(new_state, force=True)
 
@@ -422,7 +427,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             time.sleep(interval)
             state = self.set_state(new_state)
 
-    def set_state(self, state: State, force: bool = False) -> State:
+    def set_state(self, state: State[R], force: bool = False) -> State[R]:
         last_state = self.state
         if not self.task_run:
             raise ValueError("Task run is not set")
@@ -537,7 +542,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                "Task run failed with exception: %r - " "Retry %s/%s will start %s",
+                "Task run failed with exception: %r - Retry %s/%s will start %s",
                 exc,
                 self.retries + 1,
                 self.task.retries,
@@ -545,7 +550,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
 
             self.set_state(new_state, force=True)
-            self.retries = self.retries + 1
+            self.retries: int = self.retries + 1
             return True
         elif self.retries >= self.task.retries:
             self.logger.error(
@@ -639,7 +644,9 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             stack.enter_context(ConcurrencyContextV1())
             stack.enter_context(ConcurrencyContext())
 
-            self.logger = task_run_logger(task_run=self.task_run, task=self.task)  # type: ignore
+            self.logger: "logging.Logger" = task_run_logger(
+                task_run=self.task_run, task=self.task
+            )  # type: ignore
 
             yield
 
@@ -648,7 +655,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self,
         task_run_id: Optional[UUID] = None,
         dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
-    ) -> Generator["SyncTaskRunEngine", Any, Any]:
+    ) -> Generator[Self, Any, Any]:
         """
         Enters a client context and creates a task run if needed.
         """
@@ -718,7 +725,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                     self._is_started = False
                     self._client = None
 
-    async def wait_until_ready(self):
+    async def wait_until_ready(self) -> None:
         """Waits until the scheduled time (if its the future), then enters Running."""
         if scheduled_time := self.state.state_details.scheduled_time:
             sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
@@ -825,6 +832,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
 @dataclass
 class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
+    task_run: TaskRun | None = None
     _client: Optional[PrefectClient] = None
 
     @property
@@ -866,7 +874,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
             return False
 
-    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
         task = self.task
@@ -901,7 +909,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 self.logger.info(f"Hook {hook_name!r} finished running successfully")
 
-    async def begin_run(self):
+    async def begin_run(self) -> None:
         try:
             self._resolve_parameters()
             self._set_custom_task_run_name()
@@ -1067,7 +1075,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                "Task run failed with exception: %r - " "Retry %s/%s will start %s",
+                "Task run failed with exception: %r - Retry %s/%s will start %s",
                 exc,
                 self.retries + 1,
                 self.task.retries,
@@ -1075,7 +1083,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
 
             await self.set_state(new_state, force=True)
-            self.retries = self.retries + 1
+            self.retries: int = self.retries + 1
             return True
         elif self.retries >= self.task.retries:
             self.logger.error(
@@ -1169,7 +1177,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
             stack.enter_context(ConcurrencyContext())
 
-            self.logger = task_run_logger(task_run=self.task_run, task=self.task)  # type: ignore
+            self.logger: "logging.Logger" = task_run_logger(
+                task_run=self.task_run, task=self.task
+            )  # type: ignore
 
             yield
 
@@ -1178,7 +1188,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self,
         task_run_id: Optional[UUID] = None,
         dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
-    ) -> AsyncGenerator["AsyncTaskRunEngine", Any]:
+    ) -> AsyncGenerator[Self, Any]:
         """
         Enters a client context and creates a task run if needed.
         """
@@ -1246,7 +1256,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                     self._is_started = False
                     self._client = None
 
-    async def wait_until_ready(self):
+    async def wait_until_ready(self) -> None:
         """Waits until the scheduled time (if its the future), then enters Running."""
         if scheduled_time := self.state.state_details.scheduled_time:
             sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
@@ -1341,7 +1351,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         if transaction.is_committed():
             result = transaction.read()
         else:
-            if self.task_run.tags:
+            if self.task_run and self.task_run.tags:
                 # Acquire a concurrency slot for each tag, but only if a limit
                 # matching the tag already exists.
                 async with aconcurrency(list(self.task_run.tags), self.task_run.id):
@@ -1546,7 +1556,7 @@ def run_task(
     Returns:
         The result of the task run
     """
-    kwargs = dict(
+    kwargs: dict[str, Any] = dict(
         task=task,
         task_run_id=task_run_id,
         task_run=task_run,