prefect-client 2.20.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/utilities/callables.py

@@ -5,34 +5,31 @@ Utilities for working with Python callables.
 import ast
 import importlib.util
 import inspect
+import warnings
 from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
 
 import cloudpickle
-
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-from prefect._internal.pydantic.v1_schema import has_v1_type_as_param
-
-if HAS_PYDANTIC_V2:
-    import pydantic.v1 as pydantic
-
-    from prefect._internal.pydantic.v2_schema import (
-        create_v2_schema,
-        process_v2_params,
-    )
-else:
-    import pydantic
-
+import pydantic
 from griffe import Docstring, DocstringSectionKind, Parser, parse
 from typing_extensions import Literal
 
+from prefect._internal.pydantic.v1_schema import has_v1_type_as_param
+from prefect._internal.pydantic.v2_schema import (
+    create_v2_schema,
+    process_v2_params,
+)
 from prefect.exceptions import (
+    MappingLengthMismatch,
+    MappingMissingIterable,
     ParameterBindError,
     ReservedArgumentError,
     SignatureMismatchError,
 )
 from prefect.logging.loggers import disable_logger, get_logger
+from prefect.utilities.annotations import allow_failure, quote, unmapped
+from prefect.utilities.collections import isiterable
 from prefect.utilities.importtools import safe_load_namespace
 
 logger = get_logger(__name__)
@@ -45,11 +42,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:
@@ -228,15 +237,14 @@ class ParameterSchema(pydantic.BaseModel):
     title: Literal["Parameters"] = "Parameters"
     type: Literal["object"] = "object"
     properties: Dict[str, Any] = pydantic.Field(default_factory=dict)
-    required: List[str] = None
-    definitions: Optional[Dict[str, Any]] = None
+    required: List[str] = pydantic.Field(default_factory=list)
+    definitions: Dict[str, Any] = pydantic.Field(default_factory=dict)
 
-    def dict(self, *args, **kwargs):
-        """Exclude `None` fields by default to comply with
-        the OpenAPI spec.
-        """
-        kwargs.setdefault("exclude_none", True)
-        return super().dict(*args, **kwargs)
+    def model_dump_for_openapi(self) -> Dict[str, Any]:
+        result = self.model_dump(mode="python", exclude_none=True)
+        if "required" in result and not result["required"]:
+            del result["required"]
+        return result
 
 
 def parameter_docstrings(docstring: Optional[str]) -> Dict[str, str]:
@@ -284,21 +292,31 @@ def process_v1_params(
         name = param.name
 
     type_ = Any if param.annotation is inspect._empty else param.annotation
-    field = pydantic.Field(
-        default=... if param.default is param.empty else param.default,
-        title=param.name,
-        description=docstrings.get(param.name, None),
-        alias=aliases.get(name),
-        position=position,
-    )
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=pydantic.warnings.PydanticDeprecatedSince20
+        )
+        field = pydantic.Field(
+            default=... if param.default is param.empty else param.default,
+            title=param.name,
+            description=docstrings.get(param.name, None),
+            alias=aliases.get(name),
+            position=position,
+        )
     return name, type_, field
 
 
 def create_v1_schema(name_: str, model_cfg, **model_fields):
-    model: "pydantic.BaseModel" = pydantic.create_model(
-        name_, __config__=model_cfg, **model_fields
-    )
-    return model.schema(by_alias=True)
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=pydantic.warnings.PydanticDeprecatedSince20
+        )
+
+        model: "pydantic.BaseModel" = pydantic.create_model(
+            name_, __config__=model_cfg, **model_fields
+        )
+        return model.schema(by_alias=True)
 
 
 def parameter_schema(fn: Callable) -> ParameterSchema:
@@ -381,16 +399,20 @@ def generate_parameter_schema(
     model_fields = {}
     aliases = {}
 
-    class ModelConfig:
-        arbitrary_types_allowed = True
-
-    if HAS_PYDANTIC_V2 and not has_v1_type_as_param(signature):
+    if not has_v1_type_as_param(signature):
         create_schema = create_v2_schema
         process_params = process_v2_params
+
+        config = pydantic.ConfigDict(arbitrary_types_allowed=True)
     else:
         create_schema = create_v1_schema
         process_params = process_v1_params
 
+        class ModelConfig:
+            arbitrary_types_allowed = True
+
+        config = ModelConfig
+
     for position, param in enumerate(signature.parameters.values()):
         name, type_, field = process_params(
             param, position=position, docstrings=docstrings, aliases=aliases
@@ -398,16 +420,14 @@ def generate_parameter_schema(
         # Generate a Pydantic model at each step so we can check if this parameter
         # type supports schema generation
         try:
-            create_schema(
-                "CheckParameter", model_cfg=ModelConfig, **{name: (type_, field)}
-            )
+            create_schema("CheckParameter", model_cfg=config, **{name: (type_, field)})
         except (ValueError, TypeError):
             # This field's type is not valid for schema creation, update it to `Any`
             type_ = Any
         model_fields[name] = (type_, field)
 
     # Generate the final model and schema
-    schema = create_schema("Parameters", model_cfg=ModelConfig, **model_fields)
+    schema = create_schema("Parameters", model_cfg=config, **model_fields)
     return ParameterSchema(**schema)
 
 
@@ -621,3 +641,75 @@ def _get_docstring_from_source(source_code: str, func_name: str) -> Optional[str
     ):
         return func_def.body[0].value.value
     return None
+
+
+def expand_mapping_parameters(
+    func: Callable, parameters: Dict[str, Any]
+) -> List[Dict[str, Any]]:
+    """
+    Generates a list of call parameters to be used for individual calls in a mapping
+    operation.
+
+    Args:
+        func: The function to be called
+        parameters: A dictionary of parameters with iterables to be mapped over
+
+    Returns:
+        List: A list of dictionaries to be used as parameters for each
+            call in the mapping operation
+    """
+    # Ensure that any parameters in kwargs are expanded before this check
+    parameters = explode_variadic_parameter(func, parameters)
+
+    iterable_parameters = {}
+    static_parameters = {}
+    annotated_parameters = {}
+    for key, val in parameters.items():
+        if isinstance(val, (allow_failure, quote)):
+            # Unwrap annotated parameters to determine if they are iterable
+            annotated_parameters[key] = val
+            val = val.unwrap()
+
+        if isinstance(val, unmapped):
+            static_parameters[key] = val.value
+        elif isiterable(val):
+            iterable_parameters[key] = list(val)
+        else:
+            static_parameters[key] = val
+
+    if not len(iterable_parameters):
+        raise MappingMissingIterable(
+            "No iterable parameters were received. Parameters for map must "
+            f"include at least one iterable. Parameters: {parameters}"
+        )
+
+    iterable_parameter_lengths = {
+        key: len(val) for key, val in iterable_parameters.items()
+    }
+    lengths = set(iterable_parameter_lengths.values())
+    if len(lengths) > 1:
+        raise MappingLengthMismatch(
+            "Received iterable parameters with different lengths. Parameters for map"
+            f" must all be the same length. Got lengths: {iterable_parameter_lengths}"
+        )
+
+    map_length = list(lengths)[0]
+
+    call_parameters_list = []
+    for i in range(map_length):
+        call_parameters = {key: value[i] for key, value in iterable_parameters.items()}
+        call_parameters.update({key: value for key, value in static_parameters.items()})
+
+        # Add default values for parameters; these are skipped earlier since they should
+        # not be mapped over
+        for key, value in get_parameter_defaults(func).items():
+            call_parameters.setdefault(key, value)
+
+        # Re-apply annotations to each key again
+        for key, annotation in annotated_parameters.items():
+            call_parameters[key] = annotation.rewrap(call_parameters[key])
+
+        # Collapse any previously exploded kwargs
+        call_parameters_list.append(collapse_variadic_parameters(func, call_parameters))
+
+    return call_parameters_list

prefect/utilities/collections.py

@@ -4,6 +4,8 @@ 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
 from collections.abc import Sequence
@@ -28,12 +30,7 @@ from typing import (
 )
 from unittest.mock import Mock
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    import pydantic.v1 as pydantic
-else:
-    import pydantic
+import pydantic
 
 # Quote moved to `prefect.utilities.annotations` but preserved here for compatibility
 from prefect.utilities.annotations import BaseAnnotation, Quote, quote  # noqa
@@ -224,25 +221,31 @@ class StopVisiting(BaseException):
 
 
 def visit_collection(
-    expr,
-    visit_fn: 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.
+
+    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 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.
+    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)
@@ -250,30 +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], Awaitable[Any]]): an async function that
-            will be applied to every non-collection element of expr.
-        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(
@@ -284,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):
@@ -292,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:
@@ -300,90 +315,123 @@ 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
 
-    elif typ in (list, tuple, set):
+    # --- Sequences
+
+    elif isinstance(expr, (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)})
 
-    elif isinstance(expr, pydantic.BaseModel):
-        # NOTE: This implementation *does not* traverse private attributes
-        # Pydantic does not expose extras in `__fields__` so we use `__fields_set__`
-        # as well to get all of the relevant attributes
-        # Check for presence of attrs even if they're in the field set due to pydantic#4916
-        model_fields = {
-            f for f in expr.__fields_set__.union(expr.__fields__) if hasattr(expr, f)
-        }
-        items = [visit_nested(getattr(expr, key)) for key in model_fields]
+    # --- Pydantic models
 
-        if return_data:
-            # Collect fields with aliases so reconstruction can use the correct field name
-            aliases = {
-                key: value.alias
-                for key, value in expr.__fields__.items()
-                if value.has_alias
-            }
+    elif isinstance(expr, pydantic.BaseModel):
+        typ = cast(Type[pydantic.BaseModel], typ)
 
-            model_instance = typ(
-                **{
-                    aliases.get(key) or key: value
-                    for key, value in zip(model_fields, items)
-                }
-            )
+        # when extra=allow, fields not in model_fields may be in model_fields_set
+        model_fields = expr.model_fields_set.union(expr.model_fields.keys())
 
-            # Private attributes are not included in `__fields_set__` but we do not want
-            # to drop them from the model so we restore them after constructing a new
-            # model
-            for attr in expr.__private_attributes__:
-                # Use `object.__setattr__` to avoid errors on immutable models
-                object.__setattr__(model_instance, attr, getattr(expr, attr))
+        # We may encounter a deprecated field here, but this isn't the caller's fault
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=DeprecationWarning)
 
-            # Preserve data about which fields were explicitly set on the original model
-            object.__setattr__(model_instance, "__fields_set__", expr.__fields_set__)
-            result = model_instance
-        else:
-            result = None
+            updated_data = {
+                field: visit_nested(getattr(expr, field)) for field in model_fields
+            }
 
-    else:
-        result = result if return_data else None
+        if return_data:
+            modified = any(
+                getattr(expr, field) is not updated_data[field]
+                for field in model_fields
+            )
+            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
 
-    return result
+    if return_data:
+        return result
 
 
 def remove_nested_keys(keys_to_remove: List[Hashable], obj):

prefect/utilities/dispatch.py

@@ -90,8 +90,10 @@ def get_dispatch_key(
 
 @classmethod
 def _register_subclass_of_base_type(cls, **kwargs):
-    if cls.__init_subclass_original__:
+    if hasattr(cls, "__init_subclass_original__"):
         cls.__init_subclass_original__(**kwargs)
+    elif hasattr(cls, "__pydantic_init_subclass_original__"):
+        cls.__pydantic_init_subclass_original__(**kwargs)
 
     # Do not register abstract base classes
     if abc.ABC in getattr(cls, "__bases__", []):
@@ -114,8 +116,14 @@ def register_base_type(cls: T) -> T:
         registry[base_key] = cls
 
     # Add automatic subtype registration
-    cls.__init_subclass_original__ = getattr(cls, "__init_subclass__")
-    cls.__init_subclass__ = _register_subclass_of_base_type
+    if hasattr(cls, "__pydantic_init_subclass__"):
+        cls.__pydantic_init_subclass_original__ = getattr(
+            cls, "__pydantic_init_subclass__"
+        )
+        cls.__pydantic_init_subclass__ = _register_subclass_of_base_type
+    else:
+        cls.__init_subclass_original__ = getattr(cls, "__init_subclass__")
+        cls.__init_subclass__ = _register_subclass_of_base_type
 
     return cls
 
@@ -154,17 +162,22 @@ def register_type(cls: T) -> T:
     key = get_dispatch_key(cls)
     existing_value = registry.get(key)
     if existing_value is not None and id(existing_value) != id(cls):
-        # Get line numbers for debugging
-        file = inspect.getsourcefile(cls)
-        line_number = inspect.getsourcelines(cls)[1]
-        existing_file = inspect.getsourcefile(existing_value)
-        existing_line_number = inspect.getsourcelines(existing_value)[1]
-        warnings.warn(
-            f"Type {cls.__name__!r} at {file}:{line_number} has key {key!r} that "
-            f"matches existing registered type {existing_value.__name__!r} from "
-            f"{existing_file}:{existing_line_number}. The existing type will be "
-            "overridden."
-        )
+        try:
+            # Get line numbers for debugging
+            file = inspect.getsourcefile(cls)
+            line_number = inspect.getsourcelines(cls)[1]
+            existing_file = inspect.getsourcefile(existing_value)
+            existing_line_number = inspect.getsourcelines(existing_value)[1]
+            warnings.warn(
+                f"Type {cls.__name__!r} at {file}:{line_number} has key {key!r} that "
+                f"matches existing registered type {existing_value.__name__!r} from "
+                f"{existing_file}:{existing_line_number}. The existing type will be "
+                "overridden."
+            )
+        except OSError:
+            # If we can't get the source, another actor is loading this class via eval
+            # and we shouldn't update the registry
+            return cls
 
     # Add to the registry
     registry[key] = cls