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

prefect/_internal/concurrency/threads.py

@@ -8,7 +8,9 @@ import concurrent.futures
 import itertools
 import queue
 import threading
-from typing import List, Optional
+from typing import Any, Optional
+
+from typing_extensions import TypeVar
 
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.calls import Call, Portal
@@ -16,6 +18,8 @@ from prefect._internal.concurrency.cancellation import CancelledError
 from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect._internal.concurrency.primitives import Event
 
+T = TypeVar("T", infer_variance=True)
+
 
 class WorkerThread(Portal):
     """
@@ -33,7 +37,7 @@ class WorkerThread(Portal):
         self.thread = threading.Thread(
             name=name, daemon=daemon, target=self._entrypoint
         )
-        self._queue = queue.Queue()
+        self._queue: queue.Queue[Optional[Call[Any]]] = queue.Queue()
         self._run_once: bool = run_once
         self._started: bool = False
         self._submitted_count: int = 0
@@ -42,7 +46,7 @@ class WorkerThread(Portal):
         if not daemon:
             atexit.register(self.shutdown)
 
-    def start(self):
+    def start(self) -> None:
         """
         Start the worker thread.
         """
@@ -51,7 +55,7 @@ class WorkerThread(Portal):
                 self._started = True
                 self.thread.start()
 
-    def submit(self, call: Call) -> Call:
+    def submit(self, call: Call[T]) -> Call[T]:
         if self._submitted_count > 0 and self._run_once:
             raise RuntimeError(
                 "Worker configured to only run once. A call has already been submitted."
@@ -83,7 +87,7 @@ class WorkerThread(Portal):
     def name(self) -> str:
         return self.thread.name
 
-    def _entrypoint(self):
+    def _entrypoint(self) -> None:
         """
         Entrypoint for the thread.
         """
@@ -129,12 +133,14 @@ class EventLoopThread(Portal):
         self.thread = threading.Thread(
             name=name, daemon=daemon, target=self._entrypoint
         )
-        self._ready_future = concurrent.futures.Future()
+        self._ready_future: concurrent.futures.Future[
+            bool
+        ] = concurrent.futures.Future()
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._shutdown_event: Event = Event()
         self._run_once: bool = run_once
         self._submitted_count: int = 0
-        self._on_shutdown: List[Call] = []
+        self._on_shutdown: list[Call[Any]] = []
         self._lock = threading.Lock()
 
         if not daemon:
@@ -149,7 +155,7 @@ class EventLoopThread(Portal):
                 self.thread.start()
                 self._ready_future.result()
 
-    def submit(self, call: Call) -> Call:
+    def submit(self, call: Call[T]) -> Call[T]:
         if self._loop is None:
             self.start()
 
@@ -167,6 +173,7 @@ class EventLoopThread(Portal):
             call.set_runner(self)
 
             # Submit the call to the event loop
+            assert self._loop is not None
             asyncio.run_coroutine_threadsafe(self._run_call(call), self._loop)
 
             self._submitted_count += 1
@@ -180,15 +187,16 @@ class EventLoopThread(Portal):
         Shutdown the worker thread. Does not wait for the thread to stop.
         """
         with self._lock:
-            if self._shutdown_event is None:
-                return
-
             self._shutdown_event.set()
 
     @property
     def name(self) -> str:
         return self.thread.name
 
+    @property
+    def running(self) -> bool:
+        return not self._shutdown_event.is_set()
+
     def _entrypoint(self):
         """
         Entrypoint for the thread.
@@ -218,12 +226,12 @@ class EventLoopThread(Portal):
         # Empty the list to allow calls to be garbage collected. Issue #10338.
         self._on_shutdown = []
 
-    async def _run_call(self, call: Call) -> None:
+    async def _run_call(self, call: Call[Any]) -> None:
         task = call.run()
         if task is not None:
             await task
 
-    def add_shutdown_call(self, call: Call) -> None:
+    def add_shutdown_call(self, call: Call[Any]) -> None:
         self._on_shutdown.append(call)
 
     def __enter__(self):
@@ -235,9 +243,9 @@ class EventLoopThread(Portal):
 
 
 # the GLOBAL LOOP is used for background services, like logs
-GLOBAL_LOOP: Optional[EventLoopThread] = None
+_global_loop: Optional[EventLoopThread] = None
 # the RUN SYNC LOOP is used exclusively for running async functions in a sync context via asyncutils.run_sync
-RUN_SYNC_LOOP: Optional[EventLoopThread] = None
+_run_sync_loop: Optional[EventLoopThread] = None
 
 
 def get_global_loop() -> EventLoopThread:
@@ -246,29 +254,29 @@ def get_global_loop() -> EventLoopThread:
 
     Creates a new one if there is not one available.
     """
-    global GLOBAL_LOOP
+    global _global_loop
 
     # Create a new worker on first call or if the existing worker is dead
     if (
-        GLOBAL_LOOP is None
-        or not GLOBAL_LOOP.thread.is_alive()
-        or GLOBAL_LOOP._shutdown_event.is_set()
+        _global_loop is None
+        or not _global_loop.thread.is_alive()
+        or not _global_loop.running
     ):
-        GLOBAL_LOOP = EventLoopThread(daemon=True, name="GlobalEventLoopThread")
-        GLOBAL_LOOP.start()
+        _global_loop = EventLoopThread(daemon=True, name="GlobalEventLoopThread")
+        _global_loop.start()
 
-    return GLOBAL_LOOP
+    return _global_loop
 
 
 def in_global_loop() -> bool:
     """
     Check if called from the global loop.
     """
-    if GLOBAL_LOOP is None:
+    if _global_loop is None:
         # Avoid creating a global loop if there isn't one
         return False
 
-    return get_global_loop()._loop == get_running_loop()
+    return getattr(get_global_loop(), "_loop") == get_running_loop()
 
 
 def get_run_sync_loop() -> EventLoopThread:
@@ -277,29 +285,29 @@ def get_run_sync_loop() -> EventLoopThread:
 
     Creates a new one if there is not one available.
     """
-    global RUN_SYNC_LOOP
+    global _run_sync_loop
 
     # Create a new worker on first call or if the existing worker is dead
     if (
-        RUN_SYNC_LOOP is None
-        or not RUN_SYNC_LOOP.thread.is_alive()
-        or RUN_SYNC_LOOP._shutdown_event.is_set()
+        _run_sync_loop is None
+        or not _run_sync_loop.thread.is_alive()
+        or not _run_sync_loop.running
     ):
-        RUN_SYNC_LOOP = EventLoopThread(daemon=True, name="RunSyncEventLoopThread")
-        RUN_SYNC_LOOP.start()
+        _run_sync_loop = EventLoopThread(daemon=True, name="RunSyncEventLoopThread")
+        _run_sync_loop.start()
 
-    return RUN_SYNC_LOOP
+    return _run_sync_loop
 
 
 def in_run_sync_loop() -> bool:
     """
     Check if called from the global loop.
     """
-    if RUN_SYNC_LOOP is None:
+    if _run_sync_loop is None:
         # Avoid creating a global loop if there isn't one
         return False
 
-    return get_run_sync_loop()._loop == get_running_loop()
+    return getattr(get_run_sync_loop(), "_loop") == get_running_loop()
 
 
 def wait_for_global_loop_exit(timeout: Optional[float] = None) -> None:

prefect/_internal/concurrency/waiters.py

@@ -10,7 +10,8 @@ import inspect
 import queue
 import threading
 from collections import deque
-from typing import Awaitable, Generic, List, Optional, TypeVar, Union
+from collections.abc import Awaitable
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
 from weakref import WeakKeyDictionary
 
 import anyio
@@ -24,12 +25,12 @@ T = TypeVar("T")
 
 
 # Waiters are stored in a stack for each thread
-_WAITERS_BY_THREAD: "WeakKeyDictionary[threading.Thread, deque[Waiter]]" = (
+_WAITERS_BY_THREAD: "WeakKeyDictionary[threading.Thread, deque[Waiter[Any]]]" = (
     WeakKeyDictionary()
 )
 
 
-def add_waiter_for_thread(waiter: "Waiter", thread: threading.Thread):
+def add_waiter_for_thread(waiter: "Waiter[Any]", thread: threading.Thread) -> None:
     """
     Add a waiter for a thread.
     """
@@ -62,7 +63,7 @@ class Waiter(Portal, abc.ABC, Generic[T]):
         return self._call.future.done()
 
     @abc.abstractmethod
-    def wait(self) -> Union[Awaitable[None], None]:
+    def wait(self) -> T:
         """
         Wait for the call to finish.
 
@@ -71,7 +72,7 @@ class Waiter(Portal, abc.ABC, Generic[T]):
         raise NotImplementedError()
 
     @abc.abstractmethod
-    def add_done_callback(self, callback: Call) -> Call:
+    def add_done_callback(self, callback: Call[Any]) -> None:
         """
         Schedule a call to run when the waiter is done waiting.
 
@@ -91,11 +92,11 @@ class SyncWaiter(Waiter[T]):
 
     def __init__(self, call: Call[T]) -> None:
         super().__init__(call=call)
-        self._queue: queue.Queue = queue.Queue()
-        self._done_callbacks = []
+        self._queue: queue.Queue[Optional[Call[T]]] = queue.Queue()
+        self._done_callbacks: list[Call[Any]] = []
         self._done_event = threading.Event()
 
-    def submit(self, call: Call):
+    def submit(self, call: Call[T]) -> Call[T]:
         """
         Submit a callback to execute while waiting.
         """
@@ -109,7 +110,7 @@ class SyncWaiter(Waiter[T]):
     def _handle_waiting_callbacks(self):
         logger.debug("Waiter %r watching for callbacks", self)
         while True:
-            callback: Call = self._queue.get()
+            callback = self._queue.get()
             if callback is None:
                 break
 
@@ -130,13 +131,13 @@ class SyncWaiter(Waiter[T]):
                 if callback:
                     callback.run()
 
-    def add_done_callback(self, callback: Call):
+    def add_done_callback(self, callback: Call[Any]) -> None:
         if self._done_event.is_set():
             raise RuntimeError("Cannot add done callbacks to done waiters.")
         else:
             self._done_callbacks.append(callback)
 
-    def wait(self) -> T:
+    def wait(self) -> Call[T]:
         # Stop watching for work once the future is done
         self._call.future.add_done_callback(lambda _: self._queue.put_nowait(None))
         self._call.future.add_done_callback(lambda _: self._done_event.set())
@@ -159,13 +160,13 @@ class AsyncWaiter(Waiter[T]):
 
         # Delay instantiating loop and queue as there may not be a loop present yet
         self._loop: Optional[asyncio.AbstractEventLoop] = None
-        self._queue: Optional[asyncio.Queue] = None
-        self._early_submissions: List[Call] = []
-        self._done_callbacks = []
+        self._queue: Optional[asyncio.Queue[Optional[Call[T]]]] = None
+        self._early_submissions: list[Call[T]] = []
+        self._done_callbacks: list[Call[Any]] = []
         self._done_event = Event()
         self._done_waiting = False
 
-    def submit(self, call: Call):
+    def submit(self, call: Call[T]) -> Call[T]:
         """
         Submit a callback to execute while waiting.
         """
@@ -180,11 +181,15 @@ class AsyncWaiter(Waiter[T]):
             return call
 
         # We must put items in the queue from the event loop that owns it
+        if TYPE_CHECKING:
+            assert self._loop is not None
         call_soon_in_loop(self._loop, self._queue.put_nowait, call)
         return call
 
-    def _resubmit_early_submissions(self):
-        assert self._queue
+    def _resubmit_early_submissions(self) -> None:
+        if TYPE_CHECKING:
+            assert self._queue is not None
+            assert self._loop is not None
         for call in self._early_submissions:
             # We must put items in the queue from the event loop that owns it
             call_soon_in_loop(self._loop, self._queue.put_nowait, call)
@@ -192,11 +197,11 @@ class AsyncWaiter(Waiter[T]):
 
     async def _handle_waiting_callbacks(self):
         logger.debug("Waiter %r watching for callbacks", self)
-        tasks = []
+        tasks: list[Awaitable[None]] = []
 
         try:
             while True:
-                callback: Call = await self._queue.get()
+                callback = await self._queue.get()
                 if callback is None:
                     break
 
@@ -228,12 +233,12 @@ class AsyncWaiter(Waiter[T]):
                     with anyio.CancelScope(shield=True):
                         await self._run_done_callback(callback)
 
-    async def _run_done_callback(self, callback: Call):
+    async def _run_done_callback(self, callback: Call[Any]) -> None:
         coro = callback.run()
         if coro:
             await coro
 
-    def add_done_callback(self, callback: Call):
+    def add_done_callback(self, callback: Call[Any]) -> None:
         if self._done_event.is_set():
             raise RuntimeError("Cannot add done callbacks to done waiters.")
         else:
@@ -243,6 +248,8 @@ class AsyncWaiter(Waiter[T]):
         # Only send a `None` to the queue if the waiter is still blocked reading from
         # the queue. Otherwise, it's possible that the event loop is stopped.
         if not self._done_waiting:
+            assert self._loop is not None
+            assert self._queue is not None
             call_soon_in_loop(self._loop, self._queue.put_nowait, None)
 
     async def wait(self) -> Call[T]:

prefect/_internal/pydantic/v1_schema.py

@@ -6,7 +6,7 @@ import pydantic
 from pydantic.v1 import BaseModel as V1BaseModel
 
 
-def is_v1_model(v) -> bool:
+def is_v1_model(v: typing.Any) -> bool:
     with warnings.catch_warnings():
         warnings.filterwarnings(
             "ignore", category=pydantic.warnings.PydanticDeprecatedSince20
@@ -23,7 +23,7 @@ def is_v1_model(v) -> bool:
         return False
 
 
-def is_v1_type(v) -> bool:
+def is_v1_type(v: typing.Any) -> bool:
     if is_v1_model(v):
         return True
 

prefect/_internal/pydantic/v2_schema.py

@@ -16,7 +16,7 @@ from prefect._internal.pydantic.annotations.pendulum import (
 from prefect._internal.pydantic.schemas import GenerateEmptySchemaForUserClasses
 
 
-def is_v2_model(v) -> bool:
+def is_v2_model(v: t.Any) -> bool:
     if isinstance(v, V2BaseModel):
         return True
     try:
@@ -28,7 +28,7 @@ def is_v2_model(v) -> bool:
     return False
 
 
-def is_v2_type(v) -> bool:
+def is_v2_type(v: t.Any) -> bool:
     if is_v2_model(v):
         return True
 
@@ -56,9 +56,9 @@ def process_v2_params(
     param: inspect.Parameter,
     *,
     position: int,
-    docstrings: t.Dict[str, str],
-    aliases: t.Dict,
-) -> t.Tuple[str, t.Any, "pydantic.Field"]:
+    docstrings: dict[str, str],
+    aliases: dict[str, str],
+) -> tuple[str, t.Any, t.Any]:
     """
     Generate a sanitized name, type, and pydantic.Field for a given parameter.
 
@@ -72,7 +72,7 @@ def process_v2_params(
     else:
         name = param.name
 
-    type_ = t.Any if param.annotation is inspect._empty else param.annotation
+    type_ = t.Any if param.annotation is inspect.Parameter.empty else param.annotation
 
     # Replace pendulum type annotations with our own so that they are pydantic compatible
     if type_ == pendulum.DateTime:
@@ -95,12 +95,13 @@ def process_v2_params(
 def create_v2_schema(
     name_: str,
     model_cfg: t.Optional[ConfigDict] = None,
-    model_base: t.Optional[t.Type[V2BaseModel]] = None,
-    **model_fields,
-):
+    model_base: t.Optional[type[V2BaseModel]] = None,
+    model_fields: t.Optional[dict[str, t.Any]] = None,
+) -> dict[str, t.Any]:
     """
     Create a pydantic v2 model and craft a v1 compatible schema from it.
     """
+    model_fields = model_fields or {}
     model = create_model(
         name_, __config__=model_cfg, __base__=model_base, **model_fields
     )

prefect/_internal/pydantic/v2_validated_func.py

@@ -6,16 +6,17 @@ Specifically it allows for us to validate v2 models used as flow/task
 arguments.
 """
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
 
 # importing directly from v2 to be able to create a v2 model
 from pydantic import BaseModel, ConfigDict, create_model, field_validator
 from pydantic.v1.decorator import ValidatedFunction
 from pydantic.v1.errors import ConfigError
 from pydantic.v1.utils import to_camel
+from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
-    ConfigType = Union[None, Type[Any], Dict[str, Any]]
+    ConfigType: TypeAlias = Union[None, type[Any], dict[str, Any]]
 
 V_POSITIONAL_ONLY_NAME = "v__positional_only"
 V_DUPLICATE_KWARGS = "v__duplicate_kwargs"
@@ -24,13 +25,17 @@ V_DUPLICATE_KWARGS = "v__duplicate_kwargs"
 class V2ValidatedFunction(ValidatedFunction):
     def create_model(
         self,
-        fields: Dict[str, Any],
+        fields: dict[str, Any],
         takes_args: bool,
         takes_kwargs: bool,
-        config: ConfigDict,
+        config: "ConfigType",
     ) -> None:
         pos_args = len(self.arg_mapping)
 
+        config = {} if config is None else config
+        if not isinstance(config, dict):
+            raise TypeError(f"config must be None or a dict, got {type(config)}")
+
         if config.get("fields") or config.get("alias_generator"):
             raise ConfigError(
                 'Setting the "fields" and "alias_generator" property on custom Config'
@@ -42,11 +47,11 @@ class V2ValidatedFunction(ValidatedFunction):
 
         # This is the key change -- inheriting the BaseModel class from v2
         class DecoratorBaseModel(BaseModel):
-            model_config = config
+            model_config: ClassVar[ConfigDict] = ConfigDict(**config)
 
             @field_validator(self.v_args_name, check_fields=False)
             @classmethod
-            def check_args(cls, v: Optional[List[Any]]) -> Optional[List[Any]]:
+            def check_args(cls, v: Optional[list[Any]]) -> Optional[list[Any]]:
                 if takes_args or v is None:
                     return v
 
@@ -58,8 +63,8 @@ class V2ValidatedFunction(ValidatedFunction):
             @field_validator(self.v_kwargs_name, check_fields=False)
             @classmethod
             def check_kwargs(
-                cls, v: Optional[Dict[str, Any]]
-            ) -> Optional[Dict[str, Any]]:
+                cls, v: Optional[dict[str, Any]]
+            ) -> Optional[dict[str, Any]]:
                 if takes_kwargs or v is None:
                     return v
 
@@ -69,7 +74,7 @@ class V2ValidatedFunction(ValidatedFunction):
 
             @field_validator(V_POSITIONAL_ONLY_NAME, check_fields=False)
             @classmethod
-            def check_positional_only(cls, v: Optional[List[str]]) -> None:
+            def check_positional_only(cls, v: Optional[list[str]]) -> None:
                 if v is None:
                     return
 
@@ -82,7 +87,7 @@ class V2ValidatedFunction(ValidatedFunction):
 
             @field_validator(V_DUPLICATE_KWARGS, check_fields=False)
             @classmethod
-            def check_duplicate_kwargs(cls, v: Optional[List[str]]) -> None:
+            def check_duplicate_kwargs(cls, v: Optional[list[str]]) -> None:
                 if v is None:
                     return
 

prefect/_internal/retries.py

@@ -1,10 +1,15 @@
 import asyncio
 from functools import wraps
-from typing import Any, Callable, Tuple, Type
+from typing import Callable, Optional, Tuple, Type, TypeVar
+
+from typing_extensions import ParamSpec
 
 from prefect._internal._logging import logger
 from prefect.utilities.math import clamped_poisson_interval
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
 
 def exponential_backoff_with_jitter(
     attempt: int, base_delay: float, max_delay: float
@@ -21,7 +26,8 @@ def retry_async_fn(
     base_delay: float = 1,
     max_delay: float = 10,
     retry_on_exceptions: Tuple[Type[Exception], ...] = (Exception,),
-):
+    operation_name: Optional[str] = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """A decorator for retrying an async function.
 
     Args:
@@ -33,23 +39,26 @@ def retry_async_fn(
         max_delay: The maximum delay to use for the last attempt.
         retry_on_exceptions: A tuple of exception types to retry on. Defaults to
             retrying on all exceptions.
+        operation_name: Optional name to use for logging the operation instead of
+            the function name. If None, uses the function name.
     """
 
-    def decorator(func):
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
         @wraps(func)
-        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            name = operation_name or func.__name__
             for attempt in range(max_attempts):
                 try:
                     return await func(*args, **kwargs)
                 except retry_on_exceptions as e:
                     if attempt == max_attempts - 1:
                         logger.exception(
-                            f"Function {func.__name__!r} failed after {max_attempts} attempts"
+                            f"Function {name!r} failed after {max_attempts} attempts"
                         )
                         raise
                     delay = backoff_strategy(attempt, base_delay, max_delay)
                     logger.warning(
-                        f"Attempt {attempt + 1} of function {func.__name__!r} failed with {type(e).__name__}. "
+                        f"Attempt {attempt + 1} of function {name!r} failed with {type(e).__name__}: {str(e)}. "
                         f"Retrying in {delay:.2f} seconds..."
                     )
                     await asyncio.sleep(delay)

prefect/_internal/schemas/bases.py

@@ -4,7 +4,7 @@ Utilities for creating and working with Prefect REST API schemas.
 
 import datetime
 import os
-from typing import Any, ClassVar, Optional, Set, TypeVar
+from typing import Any, ClassVar, Generator, Optional, Set, TypeVar, cast
 from uuid import UUID, uuid4
 
 import pendulum
@@ -13,9 +13,10 @@ from pydantic import (
     ConfigDict,
     Field,
 )
-from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import Self
 
+from prefect.types import DateTime
+
 T = TypeVar("T")
 
 
@@ -33,7 +34,7 @@ class PrefectBaseModel(BaseModel):
 
     _reset_fields: ClassVar[Set[str]] = set()
 
-    model_config = ConfigDict(
+    model_config: ClassVar[ConfigDict] = ConfigDict(
         ser_json_timedelta="float",
         defer_build=True,
         extra=(
@@ -58,7 +59,7 @@ class PrefectBaseModel(BaseModel):
         else:
             return copy_dict == other
 
-    def __rich_repr__(self):
+    def __rich_repr__(self) -> Generator[tuple[str, Any, Any], None, None]:
         # Display all of the fields in the model if they differ from the default value
         for name, field in self.model_fields.items():
             value = getattr(self, name)
@@ -71,9 +72,11 @@ class PrefectBaseModel(BaseModel):
                 and name == "timestamp"
                 and value
             ):
-                value = pendulum.instance(value).isoformat()
+                value = cast(pendulum.DateTime, pendulum.instance(value)).isoformat()
             elif isinstance(field.annotation, datetime.datetime) and value:
-                value = pendulum.instance(value).diff_for_humans()
+                value = cast(
+                    pendulum.DateTime, pendulum.instance(value)
+                ).diff_for_humans()
 
             yield name, value, field.get_default()
 
@@ -113,11 +116,11 @@ class ObjectBaseModel(IDBaseModel):
     """
 
     _reset_fields: ClassVar[Set[str]] = {"id", "created", "updated"}
-    model_config = ConfigDict(from_attributes=True)
+    model_config: ClassVar[ConfigDict] = ConfigDict(from_attributes=True)
 
     created: Optional[DateTime] = Field(default=None, repr=False)
     updated: Optional[DateTime] = Field(default=None, repr=False)
 
 
 class ActionBaseModel(PrefectBaseModel):
-    model_config: ConfigDict = ConfigDict(extra="forbid")
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")