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

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.
 

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
     )