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

prefect/_internal/concurrency/calls.py

@@ -12,18 +12,20 @@ import dataclasses
 import inspect
 import threading
 import weakref
+from collections.abc import Awaitable, Generator
 from concurrent.futures._base import (
     CANCELLED,
     CANCELLED_AND_NOTIFIED,
     FINISHED,
     RUNNING,
 )
-from typing import Any, Awaitable, Callable, Dict, Generic, Optional, Tuple, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Union
 
-from typing_extensions import ParamSpec
+from typing_extensions import ParamSpec, Self, TypeAlias, TypeVar, TypeVarTuple
 
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.cancellation import (
+    AsyncCancelScope,
     CancelledError,
     cancel_async_at,
     cancel_sync_at,
@@ -31,9 +33,13 @@ from prefect._internal.concurrency.cancellation import (
 )
 from prefect._internal.concurrency.event_loop import get_running_loop
 
-T = TypeVar("T")
+T = TypeVar("T", infer_variance=True)
+Ts = TypeVarTuple("Ts")
 P = ParamSpec("P")
 
+_SyncOrAsyncCallable: TypeAlias = Callable[P, Union[T, Awaitable[T]]]
+
+
 # Tracks the current call being executed. Note that storing the `Call`
 # object for an async call directly in the contextvar appears to create a
 # memory leak, despite the fact that we `reset` when leaving the context
@@ -41,16 +47,16 @@ P = ParamSpec("P")
 # we already have strong references to the `Call` objects in other places
 # and b) this is used for performance optimizations where we have fallback
 # behavior if this weakref is garbage collected. A fix for issue #10952.
-current_call: contextvars.ContextVar["weakref.ref[Call]"] = (  # novm
+current_call: contextvars.ContextVar["weakref.ref[Call[Any]]"] = (  # novm
     contextvars.ContextVar("current_call")
 )
 
 # Create a strong reference to tasks to prevent destruction during execution errors
-_ASYNC_TASK_REFS = set()
+_ASYNC_TASK_REFS: set[asyncio.Task[None]] = set()
 
 
 @contextlib.contextmanager
-def set_current_call(call: "Call"):
+def set_current_call(call: "Call[Any]") -> Generator[None, Any, None]:
     token = current_call.set(weakref.ref(call))
     try:
         yield
@@ -58,7 +64,7 @@ def set_current_call(call: "Call"):
         current_call.reset(token)
 
 
-class Future(concurrent.futures.Future):
+class Future(concurrent.futures.Future[T]):
     """
     Extension of `concurrent.futures.Future` with support for cancellation of running
     futures.
@@ -70,7 +76,7 @@ class Future(concurrent.futures.Future):
         super().__init__()
         self._cancel_scope = None
         self._deadline = None
-        self._cancel_callbacks = []
+        self._cancel_callbacks: list[Callable[[], None]] = []
         self._name = name
         self._timed_out = False
 
@@ -79,7 +85,7 @@ class Future(concurrent.futures.Future):
         return super().set_running_or_notify_cancel()
 
     @contextlib.contextmanager
-    def enforce_async_deadline(self):
+    def enforce_async_deadline(self) -> Generator[AsyncCancelScope]:
         with cancel_async_at(self._deadline, name=self._name) as self._cancel_scope:
             for callback in self._cancel_callbacks:
                 self._cancel_scope.add_cancel_callback(callback)
@@ -92,7 +98,7 @@ class Future(concurrent.futures.Future):
                 self._cancel_scope.add_cancel_callback(callback)
             yield self._cancel_scope
 
-    def add_cancel_callback(self, callback: Callable[[], None]):
+    def add_cancel_callback(self, callback: Callable[[], Any]) -> None:
         """
         Add a callback to be enforced on cancellation.
 
@@ -113,7 +119,7 @@ class Future(concurrent.futures.Future):
         with self._condition:
             return self._timed_out
 
-    def cancel(self):
+    def cancel(self) -> bool:
         """Cancel the future if possible.
 
         Returns True if the future was cancelled, False otherwise. A future cannot be
@@ -147,7 +153,12 @@ class Future(concurrent.futures.Future):
         self._invoke_callbacks()
         return True
 
-    def result(self, timeout=None):
+    if TYPE_CHECKING:
+
+        def __get_result(self) -> T:
+            ...
+
+    def result(self, timeout: Optional[float] = None) -> T:
         """Return the result of the call that the future represents.
 
         Args:
@@ -186,7 +197,9 @@ class Future(concurrent.futures.Future):
             # Break a reference cycle with the exception in self._exception
             self = None
 
-    def _invoke_callbacks(self):
+    _done_callbacks: list[Callable[[Self], object]]
+
+    def _invoke_callbacks(self) -> None:
         """
         Invoke our done callbacks and clean up cancel scopes and cancel
         callbacks. Fixes a memory leak that hung on to Call objects,
@@ -206,7 +219,7 @@ class Future(concurrent.futures.Future):
 
         self._cancel_callbacks = []
         if self._cancel_scope:
-            self._cancel_scope._callbacks = []
+            setattr(self._cancel_scope, "_callbacks", [])
             self._cancel_scope = None
 
 
@@ -216,16 +229,21 @@ class Call(Generic[T]):
     A deferred function call.
     """
 
-    future: Future
-    fn: Callable[..., T]
-    args: Tuple
-    kwargs: Dict[str, Any]
+    future: Future[T]
+    fn: "_SyncOrAsyncCallable[..., T]"
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
     context: contextvars.Context
-    timeout: float
+    timeout: Optional[float]
     runner: Optional["Portal"] = None
 
     @classmethod
-    def new(cls, __fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> "Call[T]":
+    def new(
+        cls,
+        __fn: _SyncOrAsyncCallable[P, T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Self:
         return cls(
             future=Future(name=getattr(__fn, "__name__", str(__fn))),
             fn=__fn,
@@ -255,7 +273,7 @@ class Call(Generic[T]):
 
         self.runner = portal
 
-    def run(self) -> Optional[Awaitable[T]]:
+    def run(self) -> Optional[Awaitable[None]]:
         """
         Execute the call and place the result on the future.
 
@@ -337,7 +355,7 @@ class Call(Generic[T]):
     def cancel(self) -> bool:
         return self.future.cancel()
 
-    def _run_sync(self):
+    def _run_sync(self) -> Optional[Awaitable[T]]:
         cancel_scope = None
         try:
             with set_current_call(self):
@@ -348,8 +366,8 @@ class Call(Generic[T]):
                         # Forget this call's arguments in order to free up any memory
                         # that may be referenced by them; after a call has happened,
                         # there's no need to keep a reference to them
-                        self.args = None
-                        self.kwargs = None
+                        with contextlib.suppress(AttributeError):
+                            del self.args, self.kwargs
 
             # Return the coroutine for async execution
             if inspect.isawaitable(result):
@@ -357,8 +375,10 @@ class Call(Generic[T]):
 
         except CancelledError:
             # Report cancellation
+            if TYPE_CHECKING:
+                assert cancel_scope is not None
             if cancel_scope.timedout():
-                self.future._timed_out = True
+                setattr(self.future, "_timed_out", True)
                 self.future.cancel()
             elif cancel_scope.cancelled():
                 self.future.cancel()
@@ -374,8 +394,8 @@ class Call(Generic[T]):
             self.future.set_result(result)  # noqa: F821
             logger.debug("Finished call %r", self)  # noqa: F821
 
-    async def _run_async(self, coro):
-        cancel_scope = None
+    async def _run_async(self, coro: Awaitable[T]) -> None:
+        cancel_scope = result = None
         try:
             with set_current_call(self):
                 with self.future.enforce_async_deadline() as cancel_scope:
@@ -385,12 +405,14 @@ class Call(Generic[T]):
                         # Forget this call's arguments in order to free up any memory
                         # that may be referenced by them; after a call has happened,
                         # there's no need to keep a reference to them
-                        self.args = None
-                        self.kwargs = None
+                        with contextlib.suppress(AttributeError):
+                            del self.args, self.kwargs
         except CancelledError:
             # Report cancellation
+            if TYPE_CHECKING:
+                assert cancel_scope is not None
             if cancel_scope.timedout():
-                self.future._timed_out = True
+                setattr(self.future, "_timed_out", True)
                 self.future.cancel()
             elif cancel_scope.cancelled():
                 self.future.cancel()
@@ -403,10 +425,11 @@ class Call(Generic[T]):
             # Prevent reference cycle in `exc`
             del self
         else:
+            # F821 ignored because Ruff gets confused about the del self above.
             self.future.set_result(result)  # noqa: F821
             logger.debug("Finished async call %r", self)  # noqa: F821
 
-    def __call__(self) -> T:
+    def __call__(self) -> Union[T, Awaitable[T]]:
         """
         Execute the call and return its result.
 
@@ -417,7 +440,7 @@ class Call(Generic[T]):
         # Return an awaitable if in an async context
         if coro is not None:
 
-            async def run_and_return_result():
+            async def run_and_return_result() -> T:
                 await coro
                 return self.result()
 
@@ -428,8 +451,9 @@ class Call(Generic[T]):
     def __repr__(self) -> str:
         name = getattr(self.fn, "__name__", str(self.fn))
 
-        args, kwargs = self.args, self.kwargs
-        if args is None or kwargs is None:
+        try:
+            args, kwargs = self.args, self.kwargs
+        except AttributeError:
             call_args = "<dropped>"
         else:
             call_args = ", ".join(
@@ -450,7 +474,7 @@ class Portal(abc.ABC):
     """
 
     @abc.abstractmethod
-    def submit(self, call: "Call") -> "Call":
+    def submit(self, call: "Call[T]") -> "Call[T]":
         """
         Submit a call to execute elsewhere.
 

prefect/_internal/concurrency/cancellation.py

@@ -12,14 +12,15 @@ import signal
 import sys
 import threading
 import time
-from typing import Callable, Dict, Optional, Type
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Callable, Optional, overload
 
 import anyio
 
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.event_loop import get_running_loop
 
-_THREAD_SHIELDS: Dict[threading.Thread, "ThreadShield"] = {}
+_THREAD_SHIELDS: dict[threading.Thread, "ThreadShield"] = {}
 _THREAD_SHIELDS_LOCK = threading.Lock()
 
 
@@ -42,14 +43,14 @@ class ThreadShield:
         # Uses the Python implementation of the RLock instead of the C implementation
         # because we need to inspect `_count` directly to check if the lock is active
         # which is needed for delayed exception raising during alarms
-        self._lock = threading._RLock()
+        self._lock = threading._RLock()  # type: ignore  # yes, we want the private version
         self._exception = None
         self._owner = owner
 
     def __enter__(self) -> None:
         self._lock.__enter__()
 
-    def __exit__(self, *exc_info):
+    def __exit__(self, *exc_info: Any):
         retval = self._lock.__exit__(*exc_info)
 
         # Raise the exception if this is the last shield to exit in the owner thread
@@ -65,14 +66,14 @@ class ThreadShield:
 
         return retval
 
-    def set_exception(self, exc: Exception):
+    def set_exception(self, exc: BaseException):
         self._exception = exc
 
     def active(self) -> bool:
         """
         Returns true if the shield is active.
         """
-        return self._lock._count > 0
+        return getattr(self._lock, "_count") > 0
 
 
 class CancelledError(asyncio.CancelledError):
@@ -82,7 +83,7 @@ class CancelledError(asyncio.CancelledError):
     pass
 
 
-def _get_thread_shield(thread) -> ThreadShield:
+def _get_thread_shield(thread: threading.Thread) -> ThreadShield:
     with _THREAD_SHIELDS_LOCK:
         if thread not in _THREAD_SHIELDS:
             _THREAD_SHIELDS[thread] = ThreadShield(thread)
@@ -139,7 +140,7 @@ class CancelScope(abc.ABC):
         self._end_time = None
         self._timeout = timeout
         self._lock = threading.Lock()
-        self._callbacks = []
+        self._callbacks: list[Callable[[], None]] = []
         super().__init__()
 
     def __enter__(self):
@@ -151,7 +152,9 @@ class CancelScope(abc.ABC):
         logger.debug("%r entered", self)
         return self
 
-    def __exit__(self, *_):
+    def __exit__(
+        self, exc_type: type[BaseException], exc_val: Exception, exc_tb: TracebackType
+    ) -> Optional[bool]:
         with self._lock:
             if not self._cancelled:
                 self._completed = True
@@ -195,7 +198,7 @@ class CancelScope(abc.ABC):
         throw the cancelled error.
         """
         with self._lock:
-            if not self.started:
+            if not self._started:
                 raise RuntimeError("Scope has not been entered.")
 
             if self._completed:
@@ -247,7 +250,6 @@ class AsyncCancelScope(CancelScope):
         self, name: Optional[str] = None, timeout: Optional[float] = None
     ) -> None:
         super().__init__(name=name, timeout=timeout)
-        self.loop = None
 
     def __enter__(self):
         self.loop = asyncio.get_running_loop()
@@ -262,7 +264,9 @@ class AsyncCancelScope(CancelScope):
 
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self, exc_type: type[BaseException], exc_val: Exception, exc_tb: TracebackType
+    ) -> bool:
         if self._anyio_scope.cancel_called:
             # Mark as cancelled
             self.cancel(throw=False)
@@ -310,7 +314,7 @@ class NullCancelScope(CancelScope):
         super().__init__(name, timeout)
         self.reason = reason or "null cancel scope"
 
-    def cancel(self):
+    def cancel(self, throw: bool = True) -> bool:
         logger.warning("%r cannot cancel %s.", self, self.reason)
         return False
 
@@ -355,7 +359,7 @@ class AlarmCancelScope(CancelScope):
 
         return self
 
-    def _sigalarm_to_error(self, *args):
+    def _sigalarm_to_error(self, *args: object) -> None:
         logger.debug("%r captured alarm raising as cancelled error", self)
         if self.cancel(throw=False):
             shield = _get_thread_shield(threading.main_thread())
@@ -365,11 +369,13 @@ class AlarmCancelScope(CancelScope):
             else:
                 raise CancelledError()
 
-    def __exit__(self, *_):
+    def __exit__(self, *_: Any) -> Optional[bool]:
         retval = super().__exit__(*_)
 
         if self.timeout is not None:
             # Restore the previous timer
+            if TYPE_CHECKING:
+                assert self._previous_timer is not None
             signal.setitimer(signal.ITIMER_REAL, *self._previous_timer)
 
         # Restore the previous signal handler
@@ -417,7 +423,7 @@ class WatcherThreadCancelScope(CancelScope):
 
         return self
 
-    def __exit__(self, *_):
+    def __exit__(self, *_: Any) -> Optional[bool]:
         retval = super().__exit__(*_)
         self._event.set()
         if self._enforcer_thread:
@@ -466,7 +472,17 @@ class WatcherThreadCancelScope(CancelScope):
         return True
 
 
-def get_deadline(timeout: Optional[float]):
+@overload
+def get_deadline(timeout: float) -> float:
+    ...
+
+
+@overload
+def get_deadline(timeout: None) -> None:
+    ...
+
+
+def get_deadline(timeout: Optional[float]) -> Optional[float]:
     """
     Compute an deadline given a timeout.
 
@@ -572,7 +588,7 @@ def cancel_sync_after(timeout: Optional[float], name: Optional[str] = None):
         yield scope
 
 
-def _send_exception_to_thread(thread: threading.Thread, exc_type: Type[BaseException]):
+def _send_exception_to_thread(thread: threading.Thread, exc_type: type[BaseException]):
     """
     Raise an exception in a thread.
 

prefect/_internal/concurrency/event_loop.py

@@ -5,7 +5,8 @@ Thread-safe utilities for working with asynchronous event loops.
 import asyncio
 import concurrent.futures
 import functools
-from typing import Awaitable, Callable, Coroutine, Optional, TypeVar
+from collections.abc import Coroutine
+from typing import Any, Callable, Optional, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -13,7 +14,7 @@ P = ParamSpec("P")
 T = TypeVar("T")
 
 
-def get_running_loop() -> Optional[asyncio.BaseEventLoop]:
+def get_running_loop() -> Optional[asyncio.AbstractEventLoop]:
     """
     Get the current running loop.
 
@@ -30,7 +31,7 @@ def call_soon_in_loop(
     __fn: Callable[P, T],
     *args: P.args,
     **kwargs: P.kwargs,
-) -> concurrent.futures.Future:
+) -> concurrent.futures.Future[T]:
     """
     Run a synchronous call in an event loop's thread from another thread.
 
@@ -38,7 +39,7 @@ def call_soon_in_loop(
 
     Returns a future that can be used to retrieve the result of the call.
     """
-    future = concurrent.futures.Future()
+    future: concurrent.futures.Future[T] = concurrent.futures.Future()
 
     @functools.wraps(__fn)
     def wrapper() -> None:
@@ -62,8 +63,8 @@ def call_soon_in_loop(
 
 
 async def run_coroutine_in_loop_from_async(
-    __loop: asyncio.AbstractEventLoop, __coro: Coroutine
-) -> Awaitable:
+    __loop: asyncio.AbstractEventLoop, __coro: Coroutine[Any, Any, T]
+) -> T:
     """
     Run an asynchronous call in an event loop from an asynchronous context.