prefect-client 2.19.2__py3-none-any.whl → 3.0.0rc1__py3-none-any.whl

prefect/utilities/callables.py

@@ -2,36 +2,39 @@
 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.dataclasses import Docstring
 from griffe.docstrings.dataclasses import DocstringSectionKind
 from griffe.docstrings.parsers import 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
+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__)
 
 
 def get_call_parameters(
@@ -224,15 +227,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]:
@@ -280,21 +282,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:
@@ -318,20 +330,77 @@ def parameter_schema(fn: Callable) -> ParameterSchema:
         # `eval_str` is not available in Python < 3.10
         signature = inspect.signature(fn)
 
-    model_fields = {}
-    aliases = {}
     docstrings = parameter_docstrings(inspect.getdoc(fn))
 
-    class ModelConfig:
-        arbitrary_types_allowed = True
+    return generate_parameter_schema(signature, docstrings)
+
+
+def parameter_schema_from_entrypoint(entrypoint: str) -> ParameterSchema:
+    """
+    Generate a parameter schema from an entrypoint string.
+
+    Will load the source code of the function and extract the signature and docstring
+    to generate the schema.
 
-    if HAS_PYDANTIC_V2 and not has_v1_type_as_param(signature):
+    Useful for generating a schema for a function when instantiating the function may
+    not be possible due to missing imports or other issues.
+
+    Args:
+        entrypoint: A string representing the entrypoint to a function. The string
+            should be in the format of `module.path.to.function:do_stuff`.
+
+    Returns:
+        ParameterSchema: The parameter schema for the function.
+    """
+    if ":" in entrypoint:
+        # split by the last colon once to handle Windows paths with drive letters i.e C:\path\to\file.py:do_stuff
+        path, func_name = entrypoint.rsplit(":", maxsplit=1)
+        source_code = Path(path).read_text()
+    else:
+        path, func_name = entrypoint.rsplit(".", maxsplit=1)
+        spec = importlib.util.find_spec(path)
+        if not spec or not spec.origin:
+            raise ValueError(f"Could not find module {path!r}")
+        source_code = Path(spec.origin).read_text()
+    signature = _generate_signature_from_source(source_code, func_name)
+    docstring = _get_docstring_from_source(source_code, func_name)
+    return generate_parameter_schema(signature, parameter_docstrings(docstring))
+
+
+def generate_parameter_schema(
+    signature: inspect.Signature, docstrings: Dict[str, str]
+) -> ParameterSchema:
+    """
+    Generate a parameter schema from a function signature and docstrings.
+
+    To get a signature from a function, use `inspect.signature(fn)` or
+    `_generate_signature_from_source(source_code, func_name)`.
+
+    Args:
+        signature: The function signature.
+        docstrings: A dictionary mapping parameter names to docstrings.
+
+    Returns:
+        ParameterSchema: The parameter schema.
+    """
+
+    model_fields = {}
+    aliases = {}
+
+    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
@@ -339,16 +408,14 @@ def parameter_schema(fn: Callable) -> ParameterSchema:
         # 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)
 
 
@@ -362,3 +429,204 @@ def raise_for_reserved_arguments(fn: Callable, reserved_arguments: Iterable[str]
             raise ReservedArgumentError(
                 f"{argument!r} is a reserved argument name and cannot be used."
             )
+
+
+def _generate_signature_from_source(
+    source_code: str, func_name: str
+) -> inspect.Signature:
+    """
+    Extract the signature of a function from its source code.
+
+    Will ignore missing imports and exceptions while loading local class definitions.
+
+    Args:
+        source_code: The source code where the function named `func_name` is declared.
+        func_name: The name of the function.
+
+    Returns:
+        The signature of the function.
+    """
+    # Load the namespace from the source code. Missing imports and exceptions while
+    # loading local class definitions are ignored.
+    namespace = safe_load_namespace(source_code)
+    # Parse the source code into an AST
+    parsed_code = ast.parse(source_code)
+
+    func_def = next(
+        (
+            node
+            for node in ast.walk(parsed_code)
+            if isinstance(node, ast.FunctionDef) and node.name == func_name
+        ),
+        None,
+    )
+    if func_def is None:
+        raise ValueError(f"Function {func_name} not found in source code")
+    parameters = []
+
+    for arg in func_def.args.args:
+        name = arg.arg
+        annotation = arg.annotation
+        if annotation is not None:
+            try:
+                # Compile and evaluate the annotation
+                ann_code = compile(ast.Expression(annotation), "<string>", "eval")
+                annotation = eval(ann_code, namespace)
+            except Exception as e:
+                # Don't raise an error if the annotation evaluation fails. Set the
+                # annotation to `inspect.Parameter.empty` instead which is equivalent to
+                # not having an annotation.
+                logger.debug("Failed to evaluate annotation for %s: %s", name, e)
+                annotation = inspect.Parameter.empty
+        else:
+            annotation = inspect.Parameter.empty
+
+        param = inspect.Parameter(
+            name, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=annotation
+        )
+        parameters.append(param)
+
+    defaults = [None] * (
+        len(func_def.args.args) - len(func_def.args.defaults)
+    ) + func_def.args.defaults
+    for param, default in zip(parameters, defaults):
+        if default is not None:
+            try:
+                def_code = compile(ast.Expression(default), "<string>", "eval")
+                default = eval(def_code, namespace)
+            except Exception as e:
+                logger.debug(
+                    "Failed to evaluate default value for %s: %s", param.name, e
+                )
+                default = None  # Set to None if evaluation fails
+            parameters[parameters.index(param)] = param.replace(default=default)
+
+    if func_def.args.vararg:
+        parameters.append(
+            inspect.Parameter(
+                func_def.args.vararg.arg, inspect.Parameter.VAR_POSITIONAL
+            )
+        )
+    if func_def.args.kwarg:
+        parameters.append(
+            inspect.Parameter(func_def.args.kwarg.arg, inspect.Parameter.VAR_KEYWORD)
+        )
+
+    # Handle return annotation
+    return_annotation = func_def.returns
+    if return_annotation is not None:
+        try:
+            ret_ann_code = compile(
+                ast.Expression(return_annotation), "<string>", "eval"
+            )
+            return_annotation = eval(ret_ann_code, namespace)
+        except Exception as e:
+            logger.debug("Failed to evaluate return annotation: %s", e)
+            return_annotation = inspect.Signature.empty
+
+    return inspect.Signature(parameters, return_annotation=return_annotation)
+
+
+def _get_docstring_from_source(source_code: str, func_name: str) -> Optional[str]:
+    """
+    Extract the docstring of a function from its source code.
+
+    Args:
+        source_code (str): The source code of the function.
+        func_name (str): The name of the function.
+
+    Returns:
+        The docstring of the function. If the function has no docstring, returns None.
+    """
+    parsed_code = ast.parse(source_code)
+
+    func_def = next(
+        (
+            node
+            for node in ast.walk(parsed_code)
+            if isinstance(node, ast.FunctionDef) and node.name == func_name
+        ),
+        None,
+    )
+    if func_def is None:
+        raise ValueError(f"Function {func_name} not found in source code")
+
+    if (
+        func_def.body
+        and isinstance(func_def.body[0], ast.Expr)
+        and isinstance(func_def.body[0].value, ast.Constant)
+    ):
+        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,7 @@ Utilities for extensions of and operations on Python collections.
 
 import io
 import itertools
+import warnings
 from collections import OrderedDict, defaultdict
 from collections.abc import Iterator as IteratorABC
 from collections.abc import Sequence
@@ -28,12 +29,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
@@ -225,7 +221,7 @@ class StopVisiting(BaseException):
 
 def visit_collection(
     expr,
-    visit_fn: Callable[[Any], Any],
+    visit_fn: Union[Callable[[Any, dict], Any], Callable[[Any], Any]],
     return_data: bool = False,
     max_depth: int = -1,
     context: Optional[dict] = None,
@@ -252,8 +248,10 @@ def visit_collection(
 
     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.
+        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).
@@ -343,39 +341,26 @@ def visit_collection(
         result = typ(**items) if return_data else None
 
     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]
+        typ = cast(Type[pydantic.BaseModel], typ)
 
-        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
-            }
+        # 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())
 
-            model_instance = typ(
-                **{
-                    aliases.get(key) or key: value
-                    for key, value in zip(model_fields, items)
-                }
-            )
+        # We may encounter a deprecated field here, but this isn't the caller's fault
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=DeprecationWarning)
 
-            # 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))
+            updated_data = {
+                field: visit_nested(getattr(expr, field)) for field in model_fields
+            }
 
-            # Preserve data about which fields were explicitly set on the original model
-            object.__setattr__(model_instance, "__fields_set__", expr.__fields_set__)
+        if return_data:
+            # 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 = None

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
 

prefect/utilities/dockerutils.py

@@ -27,6 +27,10 @@ import prefect
 from prefect.utilities.importtools import lazy_import
 from prefect.utilities.slugify import slugify
 
+CONTAINER_LABELS = {
+    "io.prefect.version": prefect.__version__,
+}
+
 
 def python_version_minor() -> str:
     return f"{sys.version_info.major}.{sys.version_info.minor}"