prefect-client 3.0.0rc2__py3-none-any.whl → 3.0.0rc4__py3-none-any.whl

prefect/transactions.py

@@ -7,17 +7,22 @@ from typing import (
     List,
     Optional,
     Type,
-    TypeVar,
 )
 
 from pydantic import Field
+from typing_extensions import Self
 
-from prefect.context import ContextModel
+from prefect.context import ContextModel, FlowRunContext, TaskRunContext
 from prefect.records import RecordStore
+from prefect.records.result_store import ResultFactoryStore
+from prefect.results import (
+    BaseResult,
+    ResultFactory,
+    get_or_create_default_result_storage,
+)
+from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import AutoEnum
 
-T = TypeVar("T")
-
 
 class IsolationLevel(AutoEnum):
     READ_COMMITTED = AutoEnum.auto()
@@ -54,7 +59,7 @@ class Transaction(ContextModel):
     )
     overwrite: bool = False
     _staged_value: Any = None
-    __var__ = ContextVar("transaction")
+    __var__: ContextVar = ContextVar("transaction")
 
     def is_committed(self) -> bool:
         return self.state == TransactionState.COMMITTED
@@ -92,7 +97,8 @@ class Transaction(ContextModel):
         self._token = self.__var__.set(self)
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, *exc_info):
+        exc_type, exc_val, _ = exc_info
         if not self._token:
             raise RuntimeError(
                 "Asymmetric use of context. Context exit called without an enter."
@@ -123,11 +129,19 @@ class Transaction(ContextModel):
     def begin(self):
         # currently we only support READ_COMMITTED isolation
         # i.e., no locking behavior
-        if not self.overwrite and self.store and self.store.exists(key=self.key):
+        if (
+            not self.overwrite
+            and self.store
+            and self.key
+            and self.store.exists(key=self.key)
+        ):
             self.state = TransactionState.COMMITTED
 
-    def read(self) -> dict:
-        return self.store.read(key=self.key)
+    def read(self) -> BaseResult:
+        if self.store and self.key:
+            return self.store.read(key=self.key)
+        else:
+            return {}  # TODO: Determine what this should be
 
     def reset(self) -> None:
         parent = self.get_parent()
@@ -136,8 +150,9 @@ class Transaction(ContextModel):
             # parent takes responsibility
             parent.add_child(self)
 
-        self.__var__.reset(self._token)
-        self._token = None
+        if self._token:
+            self.__var__.reset(self._token)
+            self._token = None
 
         # do this below reset so that get_transaction() returns the relevant txn
         if parent and self.state == TransactionState.ROLLED_BACK:
@@ -165,7 +180,7 @@ class Transaction(ContextModel):
             for hook in self.on_commit_hooks:
                 hook(self)
 
-            if self.store:
+            if self.store and self.key:
                 self.store.write(key=self.key, value=self._staged_value)
             self.state = TransactionState.COMMITTED
             return True
@@ -174,11 +189,17 @@ class Transaction(ContextModel):
             return False
 
     def stage(
-        self, value: dict, on_rollback_hooks: list, on_commit_hooks: list
+        self,
+        value: BaseResult,
+        on_rollback_hooks: Optional[List] = None,
+        on_commit_hooks: Optional[List] = None,
     ) -> None:
         """
         Stage a value to be committed later.
         """
+        on_commit_hooks = on_commit_hooks or []
+        on_rollback_hooks = on_rollback_hooks or []
+
         if self.state != TransactionState.COMMITTED:
             self._staged_value = value
             self.on_rollback_hooks += on_rollback_hooks
@@ -203,11 +224,11 @@ class Transaction(ContextModel):
             return False
 
     @classmethod
-    def get_active(cls: Type[T]) -> Optional[T]:
+    def get_active(cls: Type[Self]) -> Optional[Self]:
         return cls.__var__.get(None)
 
 
-def get_transaction() -> Transaction:
+def get_transaction() -> Optional[Transaction]:
     return Transaction.get_active()
 
 
@@ -218,6 +239,55 @@ def transaction(
     commit_mode: CommitMode = CommitMode.LAZY,
     overwrite: bool = False,
 ) -> Generator[Transaction, None, None]:
+    """
+    A context manager for opening and managing a transaction.
+
+    Args:
+        - key: An identifier to use for the transaction
+        - store: The store to use for persisting the transaction result. If not provided,
+            a default store will be used based on the current run context.
+        - commit_mode: The commit mode controlling when the transaction and
+            child transactions are committed
+        - overwrite: Whether to overwrite an existing transaction record in the store
+
+    Yields:
+        - Transaction: An object representing the transaction state
+    """
+    # if there is no key, we won't persist a record
+    if key and not store:
+        flow_run_context = FlowRunContext.get()
+        task_run_context = TaskRunContext.get()
+        existing_factory = getattr(task_run_context, "result_factory", None) or getattr(
+            flow_run_context, "result_factory", None
+        )
+
+        if existing_factory and existing_factory.storage_block_id:
+            new_factory = existing_factory.model_copy(
+                update={
+                    "persist_result": True,
+                }
+            )
+        else:
+            default_storage = get_or_create_default_result_storage(_sync=True)
+            if existing_factory:
+                new_factory = existing_factory.model_copy(
+                    update={
+                        "persist_result": True,
+                        "storage_block": default_storage,
+                        "storage_block_id": default_storage._block_document_id,
+                    }
+                )
+            else:
+                new_factory = run_coro_as_sync(
+                    ResultFactory.default_factory(
+                        persist_result=True,
+                        result_storage=default_storage,
+                    )
+                )
+        store = ResultFactoryStore(
+            result_factory=new_factory,
+        )
+
     with Transaction(
         key=key, store=store, commit_mode=commit_mode, overwrite=overwrite
     ) as txn:

prefect/types/__init__.py

@@ -15,12 +15,19 @@ from zoneinfo import available_timezones
 MAX_VARIABLE_NAME_LENGTH = 255
 MAX_VARIABLE_VALUE_LENGTH = 5000
 
-timezone_set = available_timezones()
-
 NonNegativeInteger = Annotated[int, Field(ge=0)]
 PositiveInteger = Annotated[int, Field(gt=0)]
 NonNegativeFloat = Annotated[float, Field(ge=0.0)]
-TimeZone = Annotated[str, Field(default="UTC", pattern="|".join(timezone_set))]
+
+TimeZone = Annotated[
+    str,
+    Field(
+        default="UTC",
+        pattern="|".join(
+            [z for z in sorted(available_timezones()) if "localtime" not in z]
+        ),
+    ),
+]
 
 
 BANNED_CHARACTERS = ["/", "%", "&", ">", "<"]

prefect/utilities/asyncutils.py

@@ -314,7 +314,7 @@ def sync_compatible(async_fn: T, force_sync: bool = False) -> T:
     """
 
     @wraps(async_fn)
-    def coroutine_wrapper(*args, _sync: bool = None, **kwargs):
+    def coroutine_wrapper(*args, _sync: Optional[bool] = None, **kwargs):
         from prefect.context import MissingContextError, get_run_context
         from prefect.settings import (
             PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT,
@@ -376,8 +376,8 @@ def sync_compatible(async_fn: T, force_sync: bool = False) -> T:
 
 
 @asynccontextmanager
-async def asyncnullcontext():
-    yield
+async def asyncnullcontext(value=None):
+    yield value
 
 
 def sync(__async_fn: Callable[P, Awaitable[T]], *args: P.args, **kwargs: P.kwargs) -> T:

prefect/utilities/callables.py

@@ -44,11 +44,23 @@ def get_call_parameters(
     apply_defaults: bool = True,
 ) -> Dict[str, Any]:
     """
-    Bind a call to a function to get parameter/value mapping. Default values on the
-    signature will be included if not overridden.
-
-    Raises a ParameterBindError if the arguments/kwargs are not valid for the function
+    Bind a call to a function to get parameter/value mapping. Default values on
+    the signature will be included if not overridden.
+
+    If the function has a `__prefect_self__` attribute, it will be included as
+    the first parameter. This attribute is set when Prefect decorates a bound
+    method, so this approach allows Prefect to work with bound methods in a way
+    that is consistent with how Python handles them (i.e. users don't have to
+    pass the instance argument to the method) while still making the implicit self
+    argument visible to all of Prefect's parameter machinery (such as cache key
+    functions).
+
+    Raises a ParameterBindError if the arguments/kwargs are not valid for the
+    function
     """
+    if hasattr(fn, "__prefect_self__"):
+        call_args = (fn.__prefect_self__,) + call_args
+
     try:
         bound_signature = inspect.signature(fn).bind(*call_args, **call_kwargs)
     except TypeError as exc:

prefect/utilities/collections.py

@@ -4,6 +4,7 @@ Utilities for extensions of and operations on Python collections.
 
 import io
 import itertools
+import types
 import warnings
 from collections import OrderedDict, defaultdict
 from collections.abc import Iterator as IteratorABC
@@ -220,25 +221,31 @@ class StopVisiting(BaseException):
 
 
 def visit_collection(
-    expr,
-    visit_fn: Union[Callable[[Any, dict], Any], Callable[[Any], Any]],
+    expr: Any,
+    visit_fn: Union[Callable[[Any, Optional[dict]], Any], Callable[[Any], Any]],
     return_data: bool = False,
     max_depth: int = -1,
     context: Optional[dict] = None,
     remove_annotations: bool = False,
-):
+    _seen: Optional[Set[int]] = None,
+) -> Any:
     """
-    This function visits every element of an arbitrary Python collection. If an element
-    is a Python collection, it will be visited recursively. If an element is not a
-    collection, `visit_fn` will be called with the element. The return value of
-    `visit_fn` can be used to alter the element if `return_data` is set.
+    Visits and potentially transforms every element of an arbitrary Python collection.
 
-    Note that when using `return_data` a copy of each collection is created to avoid
-    mutating the original object. This may have significant performance penalties and
-    should only be used if you intend to transform the collection.
+    If an element is a Python collection, it will be visited recursively. If an element
+    is not a collection, `visit_fn` will be called with the element. The return value of
+    `visit_fn` can be used to alter the element if `return_data` is set to `True`.
+
+    Note:
+    - When `return_data` is `True`, a copy of each collection is created only if
+      `visit_fn` modifies an element within that collection. This approach minimizes
+      performance penalties by avoiding unnecessary copying.
+    - When `return_data` is `False`, no copies are created, and only side effects from
+      `visit_fn` are applied. This mode is faster and should be used when no transformation
+      of the collection is required, because it never has to copy any data.
 
     Supported types:
-    - List
+    - List (including iterators)
     - Tuple
     - Set
     - Dict (note: keys are also visited recursively)
@@ -246,32 +253,41 @@ def visit_collection(
     - Pydantic model
     - Prefect annotations
 
+    Note that visit_collection will not consume generators or async generators, as it would prevent
+    the caller from iterating over them.
+
     Args:
-        expr (Any): a Python object or expression
-        visit_fn (Callable[[Any, Optional[dict]], Awaitable[Any]]): a function that
-            will be applied to every non-collection element of expr. The function can
-            accept one or two arguments. If two arguments are accepted, the second
-            argument will be the context dictionary.
-        return_data (bool): if `True`, a copy of `expr` containing data modified
-            by `visit_fn` will be returned. This is slower than `return_data=False`
-            (the default).
-        max_depth: Controls the depth of recursive visitation. If set to zero, no
-            recursion will occur. If set to a positive integer N, visitation will only
-            descend to N layers deep. If set to any negative integer, no limit will be
+        expr (Any): A Python object or expression.
+        visit_fn (Callable[[Any, Optional[dict]], Any] or Callable[[Any], Any]): A function
+            that will be applied to every non-collection element of `expr`. The function can
+            accept one or two arguments. If two arguments are accepted, the second argument
+            will be the context dictionary.
+        return_data (bool): If `True`, a copy of `expr` containing data modified by `visit_fn`
+            will be returned. This is slower than `return_data=False` (the default).
+        max_depth (int): Controls the depth of recursive visitation. If set to zero, no
+            recursion will occur. If set to a positive integer `N`, visitation will only
+            descend to `N` layers deep. If set to any negative integer, no limit will be
             enforced and recursion will continue until terminal items are reached. By
             default, recursion is unlimited.
-        context: An optional dictionary. If passed, the context will be sent to each
-            call to the `visit_fn`. The context can be mutated by each visitor and will
-            be available for later visits to expressions at the given depth. Values
+        context (Optional[dict]): An optional dictionary. If passed, the context will be sent
+            to each call to the `visit_fn`. The context can be mutated by each visitor and
+            will be available for later visits to expressions at the given depth. Values
             will not be available "up" a level from a given expression.
-
             The context will be automatically populated with an 'annotation' key when
-            visiting collections within a `BaseAnnotation` type. This requires the
-            caller to pass `context={}` and will not be activated by default.
-        remove_annotations: If set, annotations will be replaced by their contents. By
+            visiting collections within a `BaseAnnotation` type. This requires the caller to
+            pass `context={}` and will not be activated by default.
+        remove_annotations (bool): If set, annotations will be replaced by their contents. By
             default, annotations are preserved but their contents are visited.
+        _seen (Optional[Set[int]]): A set of object ids that have already been visited. This
+            prevents infinite recursion when visiting recursive data structures.
+
+    Returns:
+        Any: The modified collection if `return_data` is `True`, otherwise `None`.
     """
 
+    if _seen is None:
+        _seen = set()
+
     def visit_nested(expr):
         # Utility for a recursive call, preserving options and updating the depth.
         return visit_collection(
@@ -282,6 +298,7 @@ def visit_collection(
             max_depth=max_depth - 1,
             # Copy the context on nested calls so it does not "propagate up"
             context=context.copy() if context is not None else None,
+            _seen=_seen,
         )
 
     def visit_expression(expr):
@@ -290,7 +307,7 @@ def visit_collection(
         else:
             return visit_fn(expr)
 
-    # Visit every expression
+    # --- 1. Visit every expression
     try:
         result = visit_expression(expr)
     except StopVisiting:
@@ -298,47 +315,92 @@ def visit_collection(
         result = expr
 
     if return_data:
-        # Only mutate the expression while returning data, otherwise it could be null
+        # Only mutate the root expression if the user indicated we're returning data,
+        # otherwise the function could return null and we have no collection to check
         expr = result
 
-    # Then, visit every child of the expression recursively
+    # --- 2. Visit every child of the expression recursively
 
-    # If we have reached the maximum depth, do not perform any recursion
-    if max_depth == 0:
+    # If we have reached the maximum depth or we have already visited this object,
+    # return the result if we are returning data, otherwise return None
+    if max_depth == 0 or id(expr) in _seen:
         return result if return_data else None
+    else:
+        _seen.add(id(expr))
 
     # Get the expression type; treat iterators like lists
     typ = list if isinstance(expr, IteratorABC) and isiterable(expr) else type(expr)
     typ = cast(type, typ)  # mypy treats this as 'object' otherwise and complains
 
     # Then visit every item in the expression if it is a collection
-    if isinstance(expr, Mock):
+
+    # presume that the result is the original expression.
+    # in each of the following cases, we will update the result if we need to.
+    result = expr
+
+    # --- Generators
+
+    if isinstance(expr, (types.GeneratorType, types.AsyncGeneratorType)):
+        # Do not attempt to iterate over generators, as it will exhaust them
+        pass
+
+    # --- Mocks
+
+    elif isinstance(expr, Mock):
         # Do not attempt to recurse into mock objects
-        result = expr
+        pass
+
+    # --- Annotations (unmapped, quote, etc.)
 
     elif isinstance(expr, BaseAnnotation):
         if context is not None:
             context["annotation"] = expr
-        value = visit_nested(expr.unwrap())
+        unwrapped = expr.unwrap()
+        value = visit_nested(unwrapped)
 
-        if remove_annotations:
-            result = value if return_data else None
-        else:
-            result = expr.rewrap(value) if return_data else None
+        if return_data:
+            # if we are removing annotations, return the value
+            if remove_annotations:
+                result = value
+            # if the value was modified, rewrap it
+            elif value is not unwrapped:
+                result = expr.rewrap(value)
+            # otherwise return the expr
+
+    # --- Sequences
 
     elif typ in (list, tuple, set):
         items = [visit_nested(o) for o in expr]
-        result = typ(items) if return_data else None
+        if return_data:
+            modified = any(item is not orig for item, orig in zip(items, expr))
+            if modified:
+                result = typ(items)
+
+    # --- Dictionaries
 
     elif typ in (dict, OrderedDict):
         assert isinstance(expr, (dict, OrderedDict))  # typecheck assertion
         items = [(visit_nested(k), visit_nested(v)) for k, v in expr.items()]
-        result = typ(items) if return_data else None
+        if return_data:
+            modified = any(
+                k1 is not k2 or v1 is not v2
+                for (k1, v1), (k2, v2) in zip(items, expr.items())
+            )
+            if modified:
+                result = typ(items)
+
+    # --- Dataclasses
 
     elif is_dataclass(expr) and not isinstance(expr, type):
         values = [visit_nested(getattr(expr, f.name)) for f in fields(expr)]
-        items = {field.name: value for field, value in zip(fields(expr), values)}
-        result = typ(**items) if return_data else None
+        if return_data:
+            modified = any(
+                getattr(expr, f.name) is not v for f, v in zip(fields(expr), values)
+            )
+            if modified:
+                result = typ(**{f.name: v for f, v in zip(fields(expr), values)})
+
+    # --- Pydantic models
 
     elif isinstance(expr, pydantic.BaseModel):
         typ = cast(Type[pydantic.BaseModel], typ)
@@ -355,20 +417,21 @@ def visit_collection(
             }
 
         if return_data:
-            # Use construct to avoid validation and handle immutability
-            model_instance = typ.model_construct(
-                _fields_set=expr.model_fields_set, **updated_data
+            modified = any(
+                getattr(expr, field) is not updated_data[field]
+                for field in model_fields
             )
-            for private_attr in expr.__private_attributes__:
-                setattr(model_instance, private_attr, getattr(expr, private_attr))
-            result = model_instance
-        else:
-            result = None
+            if modified:
+                # Use construct to avoid validation and handle immutability
+                model_instance = typ.model_construct(
+                    _fields_set=expr.model_fields_set, **updated_data
+                )
+                for private_attr in expr.__private_attributes__:
+                    setattr(model_instance, private_attr, getattr(expr, private_attr))
+                result = model_instance
 
-    else:
-        result = result if return_data else None
-
-    return result
+    if return_data:
+        return result
 
 
 def remove_nested_keys(keys_to_remove: List[Hashable], obj):

prefect/utilities/dockerutils.py

@@ -41,7 +41,9 @@ def python_version_micro() -> str:
 
 
 def get_prefect_image_name(
-    prefect_version: str = None, python_version: str = None, flavor: str = None
+    prefect_version: Optional[str] = None,
+    python_version: Optional[str] = None,
+    flavor: Optional[str] = None,
 ) -> str:
     """
     Get the Prefect image name matching the current Prefect and Python versions.
@@ -138,7 +140,7 @@ def build_image(
     dockerfile: str = "Dockerfile",
     tag: Optional[str] = None,
     pull: bool = False,
-    platform: str = None,
+    platform: Optional[str] = None,
     stream_progress_to: Optional[TextIO] = None,
     **kwargs,
 ) -> str:
@@ -209,7 +211,7 @@ class ImageBuilder:
         self,
         base_image: str,
         base_directory: Path = None,
-        platform: str = None,
+        platform: Optional[str] = None,
         context: Path = None,
     ):
         """Create an ImageBuilder

prefect/utilities/engine.py

@@ -786,6 +786,17 @@ def resolve_to_final_result(expr, context):
         raise StopVisiting()
 
     if isinstance(expr, NewPrefectFuture):
+        upstream_task_run = context.get("current_task_run")
+        upstream_task = context.get("current_task")
+        if (
+            upstream_task
+            and upstream_task_run
+            and expr.task_run_id == upstream_task_run.id
+        ):
+            raise ValueError(
+                f"Discovered a task depending on itself. Raising to avoid a deadlock. Please inspect the inputs and dependencies of {upstream_task.name}."
+            )
+
         expr.wait()
         state = expr.state
     elif isinstance(expr, State):

prefect/utilities/filesystem.py

@@ -1,12 +1,13 @@
 """
 Utilities for working with file systems
 """
+
 import os
 import pathlib
 import threading
 from contextlib import contextmanager
 from pathlib import Path, PureWindowsPath
-from typing import Union
+from typing import Optional, Union
 
 import fsspec
 import pathspec
@@ -32,7 +33,7 @@ def create_default_ignore_file(path: str) -> bool:
 
 
 def filter_files(
-    root: str = ".", ignore_patterns: list = None, include_dirs: bool = True
+    root: str = ".", ignore_patterns: Optional[list] = None, include_dirs: bool = True
 ) -> set:
     """
     This function accepts a root directory path and a list of file patterns to ignore, and returns
@@ -40,9 +41,7 @@ def filter_files(
 
     The specification matches that of [.gitignore files](https://git-scm.com/docs/gitignore).
     """
-    if ignore_patterns is None:
-        ignore_patterns = []
-    spec = pathspec.PathSpec.from_lines("gitwildmatch", ignore_patterns)
+    spec = pathspec.PathSpec.from_lines("gitwildmatch", ignore_patterns or [])
     ignored_files = {p.path for p in spec.match_tree_entries(root)}
     if include_dirs:
         all_files = {p.path for p in pathspec.util.iter_tree_entries(root)}

prefect/utilities/importtools.py

@@ -380,6 +380,15 @@ def safe_load_namespace(source_code: str):
 
     namespace = {"__name__": "prefect_safe_namespace_loader"}
 
+    # Remove the body of the if __name__ == "__main__": block from the AST to prevent
+    # execution of guarded code
+    new_body = []
+    for node in parsed_code.body:
+        if _is_main_block(node):
+            continue
+        new_body.append(node)
+    parsed_code.body = new_body
+
     # Walk through the AST and find all import statements
     for node in ast.walk(parsed_code):
         if isinstance(node, ast.Import):
@@ -426,3 +435,23 @@ def safe_load_namespace(source_code: str):
             except Exception as e:
                 logger.debug("Failed to compile: %s", e)
     return namespace
+
+
+def _is_main_block(node: ast.AST):
+    """
+    Check if the node is an `if __name__ == "__main__":` block.
+    """
+    if isinstance(node, ast.If):
+        try:
+            # Check if the condition is `if __name__ == "__main__":`
+            if (
+                isinstance(node.test, ast.Compare)
+                and isinstance(node.test.left, ast.Name)
+                and node.test.left.id == "__name__"
+                and isinstance(node.test.comparators[0], ast.Constant)
+                and node.test.comparators[0].value == "__main__"
+            ):
+                return True
+        except AttributeError:
+            pass
+    return False

prefect/utilities/services.py

@@ -2,7 +2,7 @@ import sys
 from collections import deque
 from traceback import format_exception
 from types import TracebackType
-from typing import Callable, Coroutine, Deque, Tuple
+from typing import Callable, Coroutine, Deque, Optional, Tuple
 
 import anyio
 import httpx
@@ -22,7 +22,7 @@ async def critical_service_loop(
     backoff: int = 1,
     printer: Callable[..., None] = print,
     run_once: bool = False,
-    jitter_range: float = None,
+    jitter_range: Optional[float] = None,
 ):
     """
     Runs the given `workload` function on the specified `interval`, while being