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

prefect/_internal/pydantic/utilities/field_validator.py

@@ -0,0 +1,135 @@
+"""
+Conditional decorator for fields depending on Pydantic version.
+"""
+
+import functools
+from inspect import signature
+from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, TypeVar, Union
+
+from typing_extensions import TypeAlias
+
+from prefect._internal.pydantic._flags import HAS_PYDANTIC_V2, USE_V2_MODELS
+
+FieldValidatorModes: TypeAlias = Literal["before", "after", "wrap", "plain"]
+T = TypeVar("T", bound=Callable[..., Any])
+
+if TYPE_CHECKING:
+    from prefect._internal.pydantic._compat import BaseModel
+
+
+def field_validator(
+    field: str,
+    /,
+    *fields: str,
+    mode: FieldValidatorModes = "after",  # v2 only
+    check_fields: Union[bool, None] = None,
+    pre: bool = False,  # v1 only
+    allow_reuse: Optional[bool] = None,
+    always: bool = False,  # v1 only
+) -> Callable[[Any], Any]:
+    """Usage docs: https://docs.pydantic.dev/2.7/concepts/validators/#field-validators
+    Returns a decorator that conditionally applies Pydantic's `field_validator` or `validator`,
+    based on the Pydantic version available, for specified field(s) of a Pydantic model.
+
+    In Pydantic V2, it uses `field_validator` allowing more granular control over validation,
+    including pre-validation and post-validation modes. In Pydantic V1, it falls back to
+    using `validator`, which is less flexible but maintains backward compatibility.
+
+    Decorate methods on the class indicating that they should be used to validate fields.
+
+    Example usage:
+    ```py
+    from typing import Any
+
+    from pydantic import (
+        BaseModel,
+        ValidationError,
+        field_validator,
+    )
+
+    class Model(BaseModel):
+        a: str
+
+        @field_validator('a')
+        @classmethod
+        def ensure_foobar(cls, v: Any):
+            if 'foobar' not in v:
+                raise ValueError('"foobar" not found in a')
+            return v
+
+    print(repr(Model(a='this is foobar good')))
+    #> Model(a='this is foobar good')
+
+    try:
+        Model(a='snap')
+    except ValidationError as exc_info:
+        print(exc_info)
+        '''
+        1 validation error for Model
+        a
+        Value error, "foobar" not found in a [type=value_error, input_value='snap', input_type=str]
+        '''
+    ```
+
+        For more in depth examples, see https://docs.pydantic.dev/latest/concepts/validators/#field-validators
+
+    Args:
+        field: The first field the `field_validator` should be called on; this is separate
+            from `fields` to ensure an error is raised if you don't pass at least one.
+        *fields: Additional field(s) the `field_validator` should be called on.
+        mode: Specifies whether to validate the fields before or after validation.
+        check_fields: Whether to check that the fields actually exist on the model.
+
+    Returns:
+        A decorator that can be used to decorate a function to be used as a field_validator.
+
+    Raises:
+        PydanticUserError:
+            - If `@field_validator` is used bare (with no fields).
+            - If the args passed to `@field_validator` as fields are not strings.
+            - If `@field_validator` applied to instance methods.
+    """
+
+    def decorator(validate_func: T) -> T:
+        if USE_V2_MODELS:
+            from pydantic import field_validator  # type: ignore
+
+            return field_validator(
+                field, *fields, mode=mode, check_fields=check_fields
+            )(validate_func)
+        elif HAS_PYDANTIC_V2:
+            from pydantic.v1 import validator  # type: ignore
+        else:
+            from pydantic import validator
+
+        # Extract the parameters of the validate_func function
+        # e.g. if validate_func has a signature of (cls, v, values, config), we want to
+        # filter the kwargs to include only those expected by validate_func, which may
+        # look like (cls, v) or (cls, v, values) etc.
+        validate_func_params = signature(validate_func).parameters
+
+        @functools.wraps(validate_func)
+        def wrapper(
+            cls: "BaseModel",
+            v: Any,
+            **kwargs: Any,
+        ) -> Any:
+            filtered_kwargs: Dict[str, Any] = {
+                k: v for k, v in kwargs.items() if k in validate_func_params
+            }
+
+            return validate_func(cls, v, **filtered_kwargs)
+
+        # In Pydantic V1, `allow_reuse` is by default False, while in Pydantic V2, it is by default True.
+        # We default to False in Pydantic V1 to maintain backward compatibility
+        # e.g. One uses @validator("a", pre=True, allow_reuse=True) in Pydantic V1
+        validator_kwargs: Dict[str, Any] = {
+            "pre": pre,
+            "always": always,
+            "check_fields": check_fields if check_fields is not None else True,
+            "allow_reuse": allow_reuse if allow_reuse is not None else False,
+        }
+
+        return validator(field, *fields, **validator_kwargs)(wrapper)  # type: ignore
+
+    return decorator

prefect/_internal/pydantic/utilities/model_construct.py

@@ -0,0 +1,56 @@
+import typing
+
+from typing_extensions import Self
+
+from prefect._internal.pydantic._base_model import BaseModel
+from prefect._internal.pydantic._flags import USE_V2_MODELS
+
+T = typing.TypeVar("T", bound="BaseModel")
+
+
+def model_construct(
+    model: typing.Type[T],
+    _fields_set: typing.Optional[typing.Set[str]] = None,
+    **values: typing.Any,
+) -> T:
+    """Creates a new instance of the `model` class with validated data.
+
+    Creates a new model setting `__dict__` and `__pydantic_fields_set__` from trusted or pre-validated data.
+    Default values are respected, but no other validation is performed.
+
+    Args:
+        _fields_set: The set of field names accepted for the Model instance.
+        values: Trusted or pre-validated data dictionary.
+
+    Returns:
+        A new instance of the `model` class with validated data.
+    """
+    if USE_V2_MODELS:
+        return model.model_construct(_fields_set=_fields_set, **values)
+    else:
+        return getattr(model, "construct")(**values)
+
+
+class ModelConstructMixin(BaseModel):
+    @classmethod
+    def model_construct(
+        cls: typing.Type["Self"],
+        _fields_set: typing.Optional[typing.Set[str]] = None,
+        **values: typing.Any,
+    ) -> "Self":
+        """Creates a new instance of the `model` class with validated data.
+
+        Creates a new model setting `__dict__` and `__pydantic_fields_set__` from trusted or pre-validated data.
+        Default values are respected, but no other validation is performed.
+
+        Args:
+            _fields_set: The set of field names accepted for the Model instance.
+            values: Trusted or pre-validated data dictionary.
+
+        Returns:
+            A new instance of the `model` class with validated data.
+        """
+        return model_construct(cls, _fields_set=_fields_set, **values)
+
+
+__all__ = ["model_construct", "ModelConstructMixin"]

prefect/_internal/pydantic/utilities/model_copy.py

@@ -0,0 +1,55 @@
+import typing
+
+from typing_extensions import Self
+
+from prefect._internal.pydantic._base_model import BaseModel
+from prefect._internal.pydantic._flags import USE_V2_MODELS
+
+T = typing.TypeVar("T", bound="BaseModel")
+
+
+def model_copy(  # type: ignore[no-redef]
+    model_instance: T,
+    *,
+    update: typing.Optional[typing.Dict[str, typing.Any]] = None,
+    deep: bool = False,
+) -> T:
+    """
+    Returns a copy of the model.
+
+    Args:
+        update: Values to change/add in the new model. Note: the data is not validated
+            before creating the new model. You should trust this data.
+        deep: Set to `True` to make a deep copy of the model.
+
+    Returns:
+        New model instance.
+    """
+    if USE_V2_MODELS:
+        return model_instance.model_copy(update=update, deep=deep)
+    else:
+        return getattr(model_instance, "copy")(update=update, deep=deep)
+
+
+class ModelCopyMixin(BaseModel):
+    def model_copy(
+        self: "Self",
+        *,
+        update: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        deep: bool = False,
+    ) -> "Self":
+        """
+        Returns a copy of the model.
+
+        Args:
+            update: Values to change/add in the new model. Note: the data is not validated
+                before creating the new model. You should trust this data.
+            deep: Set to `True` to make a deep copy of the model.
+
+        Returns:
+            New model instance.
+        """
+        return model_copy(self, update=update, deep=deep)
+
+
+__all__ = ["model_copy"]

prefect/_internal/pydantic/utilities/model_dump.py

@@ -0,0 +1,136 @@
+import typing
+import warnings as python_warnings
+
+from pydantic_core import from_json
+from typing_extensions import Self
+
+from prefect._internal.pydantic._base_model import BaseModel
+from prefect._internal.pydantic._flags import USE_V2_MODELS
+
+if typing.TYPE_CHECKING:
+    from prefect._internal.pydantic._types import IncEx
+
+T = typing.TypeVar("T", bound="BaseModel")
+
+
+def model_dump(  # type: ignore[no-redef]
+    model_instance: "BaseModel",
+    *,
+    mode: typing.Union[typing.Literal["json", "python"], str] = "python",
+    include: "IncEx" = None,
+    exclude: "IncEx" = None,
+    by_alias: bool = False,
+    exclude_unset: bool = False,
+    exclude_defaults: bool = False,
+    exclude_none: bool = False,
+    round_trip: bool = False,
+    warnings: bool = True,
+) -> typing.Dict[str, typing.Any]:
+    """
+    Generate a dictionary representation of the model, optionally specifying which fields to include or exclude.
+
+    Args:
+        mode: The mode in which `to_python` should run.
+            If mode is 'json', the output will only contain JSON serializable types.
+            If mode is 'python', the output may contain non-JSON-serializable Python objects.
+        include: A set of fields to include in the output.
+        exclude: A set of fields to exclude from the output.
+        context: Additional context to pass to the serializer.
+        by_alias: Whether to use the field's alias in the dictionary key if defined.
+        exclude_unset: Whether to exclude fields that have not been explicitly set.
+        exclude_defaults: Whether to exclude fields that are set to their default value.
+        exclude_none: Whether to exclude fields that have a value of `None`.
+        round_trip: If True, dumped values should be valid as input for non-idempotent types such as Json[T].
+        warnings: Whether to log warnings when invalid fields are encountered.
+        serialize_as_any: Whether to serialize fields with duck-typing serialization behavior.
+
+    Returns:
+        A dictionary representation of the model.
+    """
+    if USE_V2_MODELS:
+        return model_instance.model_dump(
+            mode=mode,
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+        )
+    else:
+        if mode == "json":
+            with python_warnings.catch_warnings():
+                python_warnings.simplefilter("ignore")
+                return from_json(
+                    model_instance.json(
+                        include=include,
+                        exclude=exclude,
+                        by_alias=by_alias,
+                        exclude_unset=exclude_unset,
+                        exclude_defaults=exclude_defaults,
+                        exclude_none=exclude_none,
+                    )
+                )
+
+        return getattr(model_instance, "dict")(
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+        )
+
+
+class ModelDumpMixin(BaseModel):
+    def model_dump(
+        self: "Self",
+        *,
+        mode: typing.Union[typing.Literal["json", "python"], str] = "python",
+        include: "IncEx" = None,
+        exclude: "IncEx" = None,
+        by_alias: bool = False,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        round_trip: bool = False,
+        warnings: bool = True,
+    ) -> typing.Dict[str, typing.Any]:
+        """
+        Generate a dictionary representation of the model, optionally specifying which fields to include or exclude.
+
+        Args:
+            mode: The mode in which `to_python` should run.
+                If mode is 'json', the output will only contain JSON serializable types.
+                If mode is 'python', the output may contain non-JSON-serializable Python objects.
+            include: A set of fields to include in the output.
+            exclude: A set of fields to exclude from the output.
+            context: Additional context to pass to the serializer.
+            by_alias: Whether to use the field's alias in the dictionary key if defined.
+            exclude_unset: Whether to exclude fields that have not been explicitly set.
+            exclude_defaults: Whether to exclude fields that are set to their default value.
+            exclude_none: Whether to exclude fields that have a value of `None`.
+            round_trip: If True, dumped values should be valid as input for non-idempotent types such as Json[T].
+            warnings: Whether to log warnings when invalid fields are encountered.
+            serialize_as_any: Whether to serialize fields with duck-typing serialization behavior.
+
+        Returns:
+            A dictionary representation of the model.
+        """
+        return model_dump(
+            self,
+            mode=mode,
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+        )
+
+
+__all__ = ["model_dump", "ModelDumpMixin"]

prefect/_internal/pydantic/utilities/model_dump_json.py

@@ -0,0 +1,112 @@
+import typing
+
+from typing_extensions import Self
+
+from prefect._internal.pydantic._base_model import BaseModel
+from prefect._internal.pydantic._flags import USE_V2_MODELS
+
+if typing.TYPE_CHECKING:
+    from prefect._internal.pydantic._types import IncEx
+
+T = typing.TypeVar("T", bound="BaseModel")
+
+
+def model_dump_json(
+    model_instance: "BaseModel",
+    *,
+    indent: typing.Optional[int] = None,
+    include: "IncEx" = None,
+    exclude: "IncEx" = None,
+    by_alias: bool = False,
+    exclude_unset: bool = False,
+    exclude_defaults: bool = False,
+    exclude_none: bool = False,
+    round_trip: bool = False,
+    warnings: bool = True,
+) -> str:
+    """
+    Generate a JSON representation of the model, optionally specifying which fields to include or exclude.
+
+    Args:
+        indent: If provided, the number of spaces to indent the JSON output.
+        include: A list of fields to include in the output.
+        exclude: A list of fields to exclude from the output.
+        by_alias: Whether to use the field's alias in the dictionary key if defined.
+        exclude_unset: Whether to exclude fields that have not been explicitly set.
+        exclude_defaults: Whether to exclude fields that are set to their default value.
+        exclude_none: Whether to exclude fields that have a value of `None`.
+        round_trip: If True, dumped values should be valid as input for non-idempotent types such as Json[T].
+        warnings: Whether to log warnings when invalid fields are encountered.
+
+    Returns:
+        A JSON representation of the model.
+    """
+    if USE_V2_MODELS:
+        return model_instance.model_dump_json(
+            indent=indent,
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+        )
+    else:
+        return getattr(model_instance, "json")(
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+        )
+
+
+class ModelDumpJsonMixin(BaseModel):
+    def model_dump_json(
+        self: "Self",
+        *,
+        indent: typing.Optional[int] = None,
+        include: "IncEx" = None,
+        exclude: "IncEx" = None,
+        by_alias: bool = False,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        round_trip: bool = False,
+        warnings: bool = True,
+    ) -> str:
+        """
+        Generate a JSON representation of the model, optionally specifying which fields to include or exclude.
+
+        Args:
+            indent: If provided, the number of spaces to indent the JSON output.
+            include: A list of fields to include in the output.
+            exclude: A list of fields to exclude from the output.
+            by_alias: Whether to use the field's alias in the dictionary key if defined.
+            exclude_unset: Whether to exclude fields that have not been explicitly set.
+            exclude_defaults: Whether to exclude fields that are set to their default value.
+            exclude_none: Whether to exclude fields that have a value of `None`.
+            round_trip: If True, dumped values should be valid as input for non-idempotent types such as Json[T].
+            warnings: Whether to log warnings when invalid fields are encountered.
+
+        Returns:
+            A JSON representation of the model.
+        """
+        return model_dump_json(
+            self,
+            indent=indent,
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+        )
+
+
+__all__ = ["model_dump_json", "ModelDumpJsonMixin"]

prefect/_internal/pydantic/utilities/model_fields.py

@@ -0,0 +1,50 @@
+import operator
+import typing
+
+from prefect._internal.pydantic._base_model import BaseModel, FieldInfo
+from prefect._internal.pydantic._flags import HAS_PYDANTIC_V2, USE_PYDANTIC_V2
+
+T = typing.TypeVar("T")
+
+if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
+
+    class ModelFieldMixin(BaseModel):  # type: ignore [no-redef]
+        pass
+
+else:
+
+    def cast_model_field_to_field_info(model_field: typing.Any) -> FieldInfo:
+        class _FieldInfo(BaseModel):
+            model_field: typing.Any
+
+            @property
+            def annotation(self) -> typing.Any:
+                return getattr(self.model_field, "outer_type_")
+
+            @property
+            def frozen(self) -> bool:
+                return not operator.attrgetter("field_info.allow_mutation")(
+                    self.model_field
+                )
+
+            @property
+            def json_schema_extra(self) -> typing.Dict[str, typing.Any]:
+                return operator.attrgetter("field_info.extra")(self.model_field)
+
+            def __getattr__(self, key: str) -> typing.Any:
+                return getattr(self.model_field, key)
+
+            def __repr__(self) -> str:
+                return repr(self.model_field)
+
+        return typing.cast(FieldInfo, _FieldInfo(model_field=model_field))
+
+    class ModelFieldMixin(BaseModel):
+        model_fields: typing.ClassVar[typing.Dict[str, FieldInfo]]
+
+        def __init_subclass__(cls, **kwargs: typing.Any) -> None:
+            cls.model_fields = {
+                field_name: cast_model_field_to_field_info(field)
+                for field_name, field in getattr(cls, "__fields__", {}).items()
+            }
+            super().__init_subclass__(**kwargs)

prefect/_internal/pydantic/utilities/model_fields_set.py

@@ -0,0 +1,29 @@
+import typing
+
+from prefect._internal.pydantic._base_model import BaseModel
+from prefect._internal.pydantic._flags import USE_V2_MODELS
+
+
+def model_fields_set(model_instance: "BaseModel") -> typing.Set[str]:
+    """
+    Returns a set of the model's fields.
+    """
+    if USE_V2_MODELS:
+        return getattr(model_instance, "__pydantic_fields_set__")
+    else:
+        return getattr(model_instance, "__fields_set__")
+
+
+class ModelFieldsSetMixin(BaseModel):
+    @property
+    def model_fields_set(self) -> typing.Set[str]:
+        """Returns the set of fields that have been explicitly set on this model instance.
+
+        Returns:
+            A set of strings representing the fields that have been set,
+                i.e. that were not filled from defaults.
+        """
+        return model_fields_set(self)
+
+
+__all__ = ["ModelFieldsSetMixin", "model_fields_set"]

prefect/_internal/pydantic/utilities/model_json_schema.py

@@ -0,0 +1,82 @@
+import typing
+
+from typing_extensions import Self
+
+from prefect._internal.pydantic._base_model import BaseModel
+from prefect._internal.pydantic._flags import USE_V2_MODELS
+
+if typing.TYPE_CHECKING:
+    from prefect._internal.pydantic._types import JsonSchemaMode
+
+from prefect._internal.pydantic._types import DEFAULT_REF_TEMPLATE
+
+T = typing.TypeVar("T", bound="BaseModel")
+
+
+def model_json_schema(
+    model: typing.Type[T],
+    by_alias: bool = True,
+    ref_template: str = DEFAULT_REF_TEMPLATE,
+    schema_generator: typing.Any = None,
+    mode: "JsonSchemaMode" = "validation",
+) -> typing.Dict[str, typing.Any]:
+    """
+    Generates a JSON schema for a model class.
+
+    Args:
+        by_alias: Whether to use attribute aliases or not.
+        ref_template: The reference template.
+        schema_generator: To override the logic used to generate the JSON schema, as a subclass of
+            `GenerateJsonSchema` with your desired modifications
+        mode: The mode in which to generate the schema.
+
+    Returns:
+        The JSON schema for the given model class.
+    """
+    if USE_V2_MODELS:
+        return model.model_json_schema(
+            by_alias=by_alias,
+            ref_template=ref_template,
+            mode=mode,
+            # We've changed the type of schema_generator of 'schema_generator' to 'typing.Any',
+            # which is will throw an error if its None. So, we've to pass it only if its not None.
+            **{"schema_generator": schema_generator} if schema_generator else {},
+        )
+    else:
+        return getattr(model, "schema")(
+            by_alias=by_alias,
+            ref_template=ref_template,
+        )
+
+
+class ModelJsonSchemaMixin(BaseModel):
+    @classmethod
+    def model_json_schema(
+        cls: typing.Type["Self"],
+        by_alias: bool = True,
+        ref_template: str = DEFAULT_REF_TEMPLATE,
+        schema_generator: typing.Any = None,
+        mode: "JsonSchemaMode" = "validation",
+    ) -> typing.Dict[str, typing.Any]:
+        """
+        Generates a JSON schema for a model class.
+
+        Args:
+            by_alias: Whether to use attribute aliases or not.
+            ref_template: The reference template.
+            schema_generator: To override the logic used to generate the JSON schema, as a subclass of
+                `GenerateJsonSchema` with your desired modifications
+            mode: The mode in which to generate the schema.
+
+        Returns:
+            The JSON schema for the given model class.
+        """
+        return model_json_schema(
+            cls,
+            by_alias=by_alias,
+            ref_template=ref_template,
+            mode=mode,
+        )
+
+
+__all__ = ["model_json_schema", "ModelJsonSchemaMixin"]