prefect-client 2.16.8__py3-none-any.whl → 2.16.9__py3-none-any.whl

prefect/_internal/compatibility/experimental.py

@@ -10,6 +10,7 @@ Warnings may also be disabled globally with the setting `PREFECT_EXPERIMENTAL_WA
 Some experimental features require opt-in to enable any usage. These require the setting
 `PREFECT_EXPERIMENTAL_ENABLE_<GROUP>` to be set or an error will be thrown on use.
 """
+
 import functools
 import warnings
 from typing import Any, Callable, Optional, Set, Type, TypeVar
@@ -24,7 +25,7 @@ else:
 from prefect.settings import PREFECT_EXPERIMENTAL_WARN, SETTING_VARIABLES, Setting
 from prefect.utilities.callables import get_call_parameters
 
-T = TypeVar("T", bound=Callable)
+T = TypeVar("T", bound=Callable[..., Any])
 M = TypeVar("M", bound=pydantic.BaseModel)
 
 
@@ -47,12 +48,6 @@ class ExperimentalWarning(Warning):
     """
 
 
-class ExperimentalError(Exception):
-    """
-    An exception related to experimental code.
-    """
-
-
 class ExperimentalFeature(ExperimentalWarning):
     """
     A warning displayed on use of an experimental feature.
@@ -61,6 +56,12 @@ class ExperimentalFeature(ExperimentalWarning):
     """
 
 
+class ExperimentalError(Exception):
+    """
+    An exception related to experimental code.
+    """
+
+
 class ExperimentalFeatureDisabled(ExperimentalError):
     """
     An error displayed on use of a disabled experimental feature that requires opt-in.
@@ -112,7 +113,7 @@ def experimental(
 
     group_warn = _warn_setting_for_group(group)
 
-    def decorator(fn: T):
+    def decorator(fn: T) -> T:
         @functools.wraps(fn)
         def wrapper(*args, **kwargs):
             if opt_in and not group_opt_in:

prefect/_internal/concurrency/api.py

@@ -1,6 +1,7 @@
 """
 Primary developer-facing API for concurrency management.
 """
+
 import abc
 import asyncio
 import concurrent.futures
@@ -18,6 +19,7 @@ from typing import (
 
 from typing_extensions import ParamSpec
 
+from prefect._internal.concurrency.calls import get_current_call
 from prefect._internal.concurrency.threads import (
     WorkerThread,
     get_global_loop,
@@ -104,40 +106,6 @@ class _base(abc.ABC):
         runner.submit(call)
         return call
 
-    @staticmethod
-    def call_soon_in_waiting_thread(
-        __call: Union[Callable[[], T], Call[T]],
-        thread: threading.Thread,
-        timeout: Optional[float] = None,
-    ) -> Call[T]:
-        """
-        Schedule a call for execution in the thread that is waiting for the current
-        call.
-
-        Returns the submitted call.
-        """
-        call = _cast_to_call(__call)
-        waiter = get_waiter_for_thread(thread)
-        if waiter is None:
-            raise RuntimeError(f"No waiter found for thread {thread}.")
-
-        call.set_timeout(timeout)
-        waiter.submit(call)
-        return call
-
-    @staticmethod
-    def call_in_waiting_thread(
-        __call: Union[Callable[[], T], Call[T]],
-        thread: threading.Thread,
-        timeout: Optional[float] = None,
-    ) -> T:
-        """
-        Run a call in the thread that is waiting for the current call.
-
-        Returns the result of the call.
-        """
-        raise NotImplementedError()
-
     @staticmethod
     def call_in_new_thread(
         __call: Union[Callable[[], T], Call[T]], timeout: Optional[float] = None
@@ -196,13 +164,20 @@ class from_async(_base):
         return call.result()
 
     @staticmethod
-    def call_in_waiting_thread(
+    def call_soon_in_waiting_thread(
         __call: Union[Callable[[], T], Call[T]],
         thread: threading.Thread,
         timeout: Optional[float] = None,
-    ) -> Awaitable[T]:
-        call = _base.call_soon_in_waiting_thread(__call, thread, timeout=timeout)
-        return call.aresult()
+    ) -> Call[T]:
+        call = _cast_to_call(__call)
+        parent_call = get_current_call()
+        waiter = get_waiter_for_thread(thread, parent_call)
+        if waiter is None:
+            raise RuntimeError(f"No waiter found for thread {thread}.")
+
+        call.set_timeout(timeout)
+        waiter.submit(call)
+        return call
 
     @staticmethod
     def call_in_new_thread(
@@ -257,13 +232,19 @@ class from_sync(_base):
         return call.result()
 
     @staticmethod
-    def call_in_waiting_thread(
+    def call_soon_in_waiting_thread(
         __call: Union[Callable[[], T], Call[T]],
         thread: threading.Thread,
         timeout: Optional[float] = None,
-    ) -> T:
-        call = _base.call_soon_in_waiting_thread(__call, thread, timeout=timeout)
-        return call.result()
+    ) -> Call[T]:
+        call = _cast_to_call(__call)
+        waiter = get_waiter_for_thread(thread)
+        if waiter is None:
+            raise RuntimeError(f"No waiter found for thread {thread}.")
+
+        call.set_timeout(timeout)
+        waiter.submit(call)
+        return call
 
     @staticmethod
     def call_in_new_thread(

prefect/_internal/concurrency/waiters.py

@@ -9,9 +9,9 @@ import contextlib
 import inspect
 import queue
 import threading
-import weakref
 from collections import deque
 from typing import Awaitable, Generic, List, Optional, TypeVar, Union
+from weakref import WeakKeyDictionary
 
 import anyio
 
@@ -24,34 +24,37 @@ T = TypeVar("T")
 
 
 # Waiters are stored in a stack for each thread
-_WAITERS_BY_THREAD: "weakref.WeakKeyDictionary[threading.Thread, deque[Waiter]]" = (
-    weakref.WeakKeyDictionary()
+_WAITERS_BY_THREAD: "WeakKeyDictionary[threading.Thread, deque[Waiter]]" = (
+    WeakKeyDictionary()
 )
 
 
-def get_waiter_for_thread(thread: threading.Thread) -> Optional["Waiter"]:
+def get_waiter_for_thread(
+    thread: threading.Thread, parent_call: Optional[Call] = None
+) -> Optional["Waiter"]:
     """
-    Get the current waiter for a thread.
+    Get the current waiter for a thread and an optional parent call.
 
-    Returns `None` if one does not exist.
+    To avoid assigning outer callbacks to inner waiters in the case of nested calls,
+    the parent call is used to determine which waiter to return. If a parent call is
+    not provided, we return the most recently created waiter (last in the stack).
+
+    see https://github.com/PrefectHQ/prefect/issues/12036
+
+    Returns `None` if no active waiter is found for the thread.
     """
-    waiters = _WAITERS_BY_THREAD.get(thread)
-
-    if waiters:
-        idx = -1
-        while abs(idx) <= len(waiters):
-            try:
-                waiter = waiters[idx]
-                if not waiter.call_is_done():
-                    return waiter
-                idx = idx - 1
-            # It is possible that items are being added or removed
-            # from the deque, so the index we're using may not always
-            # be valid.
-            except IndexError:
-                break
 
-    return None
+    waiters: "Optional[deque[Waiter]]" = _WAITERS_BY_THREAD.get(thread)
+
+    if waiters and (active_waiters := [w for w in waiters if not w.call_is_done()]):
+        if parent_call and (
+            matching_waiter := next(
+                (w for w in active_waiters if w._call == parent_call), None
+            )
+        ):  # if exists an active waiter responsible for the parent call, return it
+            return matching_waiter
+        else:  # otherwise, return the most recently created waiter
+            return active_waiters[-1]
 
 
 def add_waiter_for_thread(waiter: "Waiter", thread: threading.Thread):

prefect/_internal/pydantic/__init__.py

@@ -6,21 +6,25 @@
 ### This is a tradeoff we're willing to make for now until pydantic v1 is
 ### no longer supported.
 
-from pydantic.version import VERSION as PYDANTIC_VERSION
 
-HAS_PYDANTIC_V2 = PYDANTIC_VERSION.startswith("2.")
+from ._flags import HAS_PYDANTIC_V2
 
 from ._compat import (
     model_dump,
     model_json_schema,
     model_validate,
-    IncEx,
     model_dump_json,
     model_copy,
     model_validate_json,
+    TypeAdapter,
     validate_python,
+    BaseModel,
+    Field,
+    FieldInfo,
 )
 
+from ._types import IncEx
+
 __all__ = [
     "model_dump",
     "model_json_schema",
@@ -29,5 +33,10 @@ __all__ = [
     "model_dump_json",
     "model_copy",
     "model_validate_json",
+    "TypeAdapter",
     "validate_python",
+    "BaseModel",
+    "HAS_PYDANTIC_V2",
+    "Field",
+    "FieldInfo",
 ]

prefect/_internal/pydantic/_base_model.py

@@ -6,11 +6,14 @@ from prefect._internal.pydantic._flags import (
 )
 
 if typing.TYPE_CHECKING:
-    from pydantic import BaseModel
+    from pydantic import BaseModel, Field
+    from pydantic.fields import FieldInfo
 
 if HAS_PYDANTIC_V2 and not USE_PYDANTIC_V2:
-    from pydantic.v1 import BaseModel
+    from pydantic.v1 import BaseModel, Field
+    from pydantic.v1.fields import FieldInfo
 else:
-    from pydantic import BaseModel
+    from pydantic import BaseModel, Field
+    from pydantic.fields import FieldInfo
 
-__all__ = ["BaseModel"]
+__all__ = ["BaseModel", "Field", "FieldInfo"]