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

prefect/_internal/pydantic/__init__.py

@@ -10,4 +10,24 @@ from pydantic.version import VERSION as PYDANTIC_VERSION
 
 HAS_PYDANTIC_V2 = PYDANTIC_VERSION.startswith("2.")
 
-from ._compat import model_dump, model_json_schema, model_validate, IncEx
+from ._compat import (
+    model_dump,
+    model_json_schema,
+    model_validate,
+    IncEx,
+    model_dump_json,
+    model_copy,
+    model_validate_json,
+    validate_python,
+)
+
+__all__ = [
+    "model_dump",
+    "model_json_schema",
+    "model_validate",
+    "IncEx",
+    "model_dump_json",
+    "model_copy",
+    "model_validate_json",
+    "validate_python",
+]

prefect/_internal/pydantic/_base_model.py

@@ -0,0 +1,16 @@
+import typing
+
+from prefect._internal.pydantic._flags import (
+    HAS_PYDANTIC_V2,
+    USE_PYDANTIC_V2,
+)
+
+if typing.TYPE_CHECKING:
+    from pydantic import BaseModel
+
+if HAS_PYDANTIC_V2 and not USE_PYDANTIC_V2:
+    from pydantic.v1 import BaseModel
+else:
+    from pydantic import BaseModel
+
+__all__ = ["BaseModel"]

prefect/_internal/pydantic/_compat.py

@@ -1,86 +1,145 @@
-from typing import Any, Dict, Literal, Optional, Set, Type, Union
+from typing import Any, Dict, Generic, Literal, Optional, Set, Type, TypeVar, Union
 
-import typing_extensions
-from pydantic import BaseModel
+from typing_extensions import Self, TypeAlias
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
+from prefect._internal.pydantic._flags import (
+    HAS_PYDANTIC_V2,
+    USE_PYDANTIC_V2,
+)
 from prefect.logging.loggers import get_logger
-from prefect.settings import PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS
 
-IncEx: typing_extensions.TypeAlias = (
-    "Union[Set[int], Set[str], Dict[int, Any], Dict[str, Any], None]"
-)
+from ._base_model import BaseModel as PydanticBaseModel
 
+IncEx: TypeAlias = "Union[Set[int], Set[str], Dict[int, Any], Dict[str, Any], None]"
 logger = get_logger("prefect._internal.pydantic")
 
+T = TypeVar("T")
+B = TypeVar("B", bound=PydanticBaseModel)
+
 if HAS_PYDANTIC_V2:
-    from pydantic.json_schema import GenerateJsonSchema
+    from pydantic import (
+        TypeAdapter as BaseTypeAdapter,
+    )
+    from pydantic import (
+        parse_obj_as,  # type: ignore
+    )
+    from pydantic.json_schema import GenerateJsonSchema  # type: ignore
+else:
+    from pydantic import parse_obj_as  # type: ignore
 
+if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
+    TypeAdapter = BaseTypeAdapter  # type: ignore
 
-def is_pydantic_v2_compatible(
-    model_instance: Optional[BaseModel] = None, fn_name: Optional[str] = None
-) -> bool:
-    """
-    Determines if the current environment is compatible with Pydantic V2 features,
-    based on the presence of Pydantic V2 and a global setting that enables V2 functionalities.
-
-    This function primarily serves to facilitate conditional logic in code that needs to
-    operate differently depending on the availability of Pydantic V2 features. It checks
-    two conditions: whether Pydantic V2 is installed, and whether the use of V2 features
-    is explicitly enabled through a global setting (`PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS`).
-
-    Parameters:
-    -----------
-    model_instance : Optional[BaseModel], optional
-        An instance of a Pydantic model. This parameter is used to perform a type check
-        to ensure the passed object is a Pydantic model instance. If not provided or if
-        the object is not a Pydantic model, a TypeError is raised. Defaults to None.
-
-    fn_name : Optional[str], optional
-        The name of the function or feature for which V2 compatibility is being checked.
-        This is used for logging purposes to provide more context in debug messages.
-        Defaults to None.
+else:
 
-    Returns:
-    --------
-    bool
-        True if the current environment supports Pydantic V2 features and if the global
-        setting for enabling V2 features is set to True. False otherwise.
+    class TypeAdapter(Generic[T]):
+        def __init__(self, type_: Union[T, Type[T]]) -> None:
+            self.type_ = type_
 
-    Raises:
-    -------
-    TypeError
-        If `model_instance` is provided but is not an instance of a Pydantic BaseModel.
+        def validate_python(
+            self,
+            __object: Any,
+            /,
+            *,
+            strict: Optional[bool] = None,
+            from_attributes: Optional[bool] = None,
+            context: Optional[Dict[str, Any]] = None,
+        ) -> Any:
+            return parse_obj_as(self.type_, __object)  # type: ignore
+
+
+# BaseModel methods and definitions
+
+
+def model_copy(
+    model_instance: PydanticBaseModel,
+    *,
+    update: Optional[Dict[str, Any]] = None,
+    deep: bool = False,
+) -> PydanticBaseModel:
+    """Usage docs: https://docs.pydantic.dev/2.7/concepts/serialization/#model_copy
+
+    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 model_instance and not isinstance(model_instance, BaseModel):
-        raise TypeError(
-            f"Expected a Pydantic model, but got {type(model_instance).__name__}"
-        )
+    if not hasattr(model_instance, "copy") and not hasattr(
+        model_instance, "model_copy"
+    ):
+        raise TypeError("Expected a Pydantic model instance")
 
-    should_dump_as_v2_model = (
-        HAS_PYDANTIC_V2 and PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS
-    )
+    if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
+        return model_instance.model_copy(update=update, deep=deep)
 
-    if should_dump_as_v2_model:
-        logger.debug(
-            f"Using Pydantic v2 compatibility layer for `{fn_name}`. This will be removed in a future release."
-        )
+    return model_instance.copy(update=update, deep=deep)  # type: ignore
 
-        return True
 
-    elif HAS_PYDANTIC_V2:
-        logger.debug(
-            "Pydantic v2 compatibility layer is disabled. To enable, set `PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS` to `True`."
-        )
+def model_dump_json(
+    model_instance: PydanticBaseModel,
+    *,
+    indent: 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.
 
-    else:
-        logger.debug("Pydantic v2 is not installed.")
+    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.
 
-    return False
+    Returns:
+        A JSON representation of the model.
+    """
+    if not hasattr(model_instance, "json") and not hasattr(
+        model_instance, "model_dump_json"
+    ):
+        raise TypeError("Expected a Pydantic model instance")
+
+    if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
+        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,
+        )
+
+    return model_instance.json(  # type: ignore
+        include=include,
+        exclude=exclude,
+        by_alias=by_alias,
+        exclude_unset=exclude_unset,
+        exclude_defaults=exclude_defaults,
+        exclude_none=exclude_none,
+    )
 
 
 def model_dump(
-    model_instance: BaseModel,
+    model_instance: PydanticBaseModel,
     *,
     mode: Union[Literal["json", "python"], str] = "python",
     include: IncEx = None,
@@ -111,7 +170,12 @@ def model_dump(
     Returns:
         A dictionary representation of the model.
     """
-    if is_pydantic_v2_compatible(model_instance=model_instance, fn_name="model_dump"):
+    if not hasattr(model_instance, "dict") and not hasattr(
+        model_instance, "model_dump"
+    ):
+        raise TypeError("Expected a Pydantic model instance")
+
+    if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
         return model_instance.model_dump(
             mode=mode,
             include=include,
@@ -124,7 +188,7 @@ def model_dump(
             warnings=warnings,
         )
 
-    return model_instance.dict(
+    return getattr(model_instance, "dict")(
         include=include,
         exclude=exclude,
         by_alias=by_alias,
@@ -139,11 +203,11 @@ JsonSchemaMode = Literal["validation", "serialization"]
 
 
 def model_json_schema(
-    model: Type[BaseModel],
+    model: Type[PydanticBaseModel],
     *,
     by_alias: bool = True,
     ref_template: str = DEFAULT_REF_TEMPLATE,
-    schema_generator=None,
+    schema_generator: Any = None,
     mode: JsonSchemaMode = "validation",
 ) -> Dict[str, Any]:
     """
@@ -165,8 +229,11 @@ def model_json_schema(
     dict[str, Any]
         The JSON schema for the given model class.
     """
-    if is_pydantic_v2_compatible(fn_name="model_json_schema"):
-        schema_generator = GenerateJsonSchema
+    if not hasattr(model, "schema") and not hasattr(model, "model_json_schema"):
+        raise TypeError("Expected a Pydantic model type")
+
+    if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
+        schema_generator = GenerateJsonSchema  # type: ignore
         return model.model_json_schema(
             by_alias=by_alias,
             ref_template=ref_template,
@@ -174,20 +241,20 @@ def model_json_schema(
             mode=mode,
         )
 
-    return model.schema(
+    return model.schema(  # type: ignore
         by_alias=by_alias,
         ref_template=ref_template,
     )
 
 
 def model_validate(
-    model: Type[BaseModel],
+    model: Type[B],
     obj: Any,
     *,
-    strict: bool = False,
-    from_attributes: bool = False,
+    strict: Optional[bool] = False,
+    from_attributes: Optional[bool] = False,
     context: Optional[Dict[str, Any]] = None,
-) -> Union[BaseModel, Dict[str, Any]]:
+) -> B:
     """Validate a pydantic model instance.
 
     Args:
@@ -202,7 +269,10 @@ def model_validate(
     Returns:
         The validated model instance.
     """
-    if is_pydantic_v2_compatible(fn_name="model_validate"):
+    if not hasattr(model, "parse_obj") and not hasattr(model, "model_validate"):
+        raise TypeError("Expected a Pydantic model type")
+
+    if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
         return model.model_validate(
             obj=obj,
             strict=strict,
@@ -210,4 +280,185 @@ def model_validate(
             context=context,
         )
 
-    return model.parse_obj(obj)
+    return getattr(model, "parse_obj")(obj)
+
+
+def model_validate_json(
+    model: Type[B],
+    json_data: Union[str, bytes, bytearray],
+    *,
+    strict: Optional[bool] = False,
+    context: Optional[Dict[str, Any]] = None,
+) -> B:
+    """Validate the given JSON data against the Pydantic model.
+
+    Args:
+        json_data: The JSON data to validate.
+        strict: Whether to enforce types strictly.
+        context: Extra variables to pass to the validator.
+
+    Returns:
+        The validated Pydantic model.
+
+    Raises:
+        ValueError: If `json_data` is not a JSON string.
+    """
+    if not hasattr(model, "parse_raw") and not hasattr(model, "model_validate_json"):
+        raise TypeError("Expected a Pydantic model type")
+
+    if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
+        return model.model_validate_json(
+            json_data=json_data,
+            strict=strict,
+            context=context,
+        )
+
+    return getattr(model, "parse_raw")(json_data)
+
+
+if HAS_PYDANTIC_V2 and USE_PYDANTIC_V2:
+    # In this case, there's no functionality to add, so we just alias the Pydantic v2 BaseModel
+    class BaseModel(PydanticBaseModel):  # type: ignore
+        pass
+
+else:
+    # In this case, we're working with a Pydantic v1 model, so we need to add Pydantic v2 functionality
+    class BaseModel(PydanticBaseModel):
+        def model_dump(
+            self: "BaseModel",
+            *,
+            mode: 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,
+        ) -> Dict[str, Any]:
+            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,
+            )
+
+        def model_dump_json(
+            self,
+            *,
+            indent: Optional[int] = None,
+            include: Optional[IncEx] = None,
+            exclude: Optional[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:
+            return model_dump_json(
+                model_instance=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,
+            )
+
+        def model_copy(
+            self: "Self",
+            *,
+            update: Optional[Dict[str, Any]] = None,
+            deep: bool = False,
+        ) -> "Self":
+            return super().model_copy(update=update, deep=deep)
+
+        @classmethod
+        def model_json_schema(
+            cls,
+            by_alias: bool = True,
+            ref_template: str = DEFAULT_REF_TEMPLATE,
+            schema_generator: Any = None,
+            mode: JsonSchemaMode = "validation",
+        ) -> Dict[str, Any]:
+            return model_json_schema(
+                cls,
+                by_alias=by_alias,
+                ref_template=ref_template,
+                schema_generator=schema_generator,
+                mode=mode,
+            )
+
+        @classmethod
+        def model_validate(
+            cls: Type["Self"],
+            obj: Any,
+            *,
+            strict: Optional[bool] = False,
+            from_attributes: Optional[bool] = False,
+            context: Optional[Dict[str, Any]] = None,
+        ) -> "Self":
+            return model_validate(
+                cls,
+                obj,
+                strict=strict,
+                from_attributes=from_attributes,
+                context=context,
+            )
+
+        @classmethod
+        def model_validate_json(
+            cls: Type["Self"],
+            json_data: Union[str, bytes, bytearray],
+            *,
+            strict: Optional[bool] = False,
+            context: Optional[Dict[str, Any]] = None,
+        ) -> "Self":
+            return model_validate_json(
+                cls,
+                json_data,
+                strict=strict,
+                context=context,
+            )
+
+
+# TypeAdapter methods and definitions
+
+
+def validate_python(
+    type_: Union[T, Type[T]],
+    __object: Any,
+    /,
+    *,
+    strict: Optional[bool] = None,
+    from_attributes: Optional[bool] = None,
+    context: Optional[Dict[str, Any]] = None,
+) -> T:
+    """Validate a Python object against the model.
+
+    Args:
+        type_: The type to validate against.
+        __object: The Python object to validate against the model.
+        strict: Whether to strictly check types.
+        from_attributes: Whether to extract data from object attributes.
+        context: Additional context to pass to the validator.
+
+    !!! note
+        When using `TypeAdapter` with a Pydantic `dataclass`, the use of the `from_attributes`
+        argument is not supported.
+
+    Returns:
+        The validated object.
+    """
+    return TypeAdapter(type_).validate_python(__object)

prefect/_internal/pydantic/_flags.py

@@ -0,0 +1,15 @@
+import os
+
+# Retrieve current version of Pydantic installed in environment
+from pydantic.version import VERSION as PYDANTIC_VERSION
+
+# Check if Pydantic version 2 is the installed version
+HAS_PYDANTIC_V2 = PYDANTIC_VERSION.startswith("2.")
+
+# Determine if Pydantic v2 internals should be used based on an environment variable.
+USE_PYDANTIC_V2 = os.environ.get(
+    "PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS", False
+) in {"1", "true", "True"}
+
+# Set to True if Pydantic v2 is present but not enabled, indicating deprecation warnings may occur.
+EXPECT_DEPRECATION_WARNINGS = HAS_PYDANTIC_V2 and not USE_PYDANTIC_V2