zenml-nightly 0.58.2.dev20240615__py3-none-any.whl → 0.58.2.dev20240623__py3-none-any.whl

zenml/utils/pydantic_utils.py

@@ -15,15 +15,28 @@
 
 import inspect
 import json
+from json.decoder import JSONDecodeError
 from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union, cast
 
 import yaml
-from pydantic import BaseModel
-from pydantic.decorator import ValidatedFunction
-from pydantic.json import pydantic_encoder
-from pydantic.utils import sequence_like
-
-from zenml.utils import dict_utils, yaml_utils
+from pydantic import (
+    AfterValidator,
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    PlainValidator,
+    ValidationInfo,
+    WrapValidator,
+    validate_call,
+)
+from pydantic._internal import _repr as pydantic_repr
+from pydantic.v1.utils import sequence_like
+
+from zenml.logger import get_logger
+from zenml.utils import dict_utils, typing_utils, yaml_utils
+from zenml.utils.json_utils import pydantic_encoder
+
+logger = get_logger(__name__)
 
 M = TypeVar("M", bound="BaseModel")
 
@@ -54,9 +67,9 @@ def update_model(
         else:
             update_dict = update
     else:
-        update_dict = update.dict(exclude_unset=True)
+        update_dict = update.model_dump(exclude_unset=True)
 
-    original_dict = original.dict(exclude_unset=True)
+    original_dict = original.model_dump(exclude_unset=True)
     if recursive:
         values = dict_utils.recursive_update(original_dict, update_dict)
     else:
@@ -94,7 +107,7 @@ class TemplateGenerator:
                 self.instance_or_class
             )
 
-        # Convert to json in an intermediate step so we can leverage Pydantic's
+        # Convert to json in an intermediate step, so we can leverage Pydantic's
         # encoder to support types like UUID and datetime
         json_string = json.dumps(template, default=pydantic_encoder)
         return cast(Dict[str, Any], json.loads(json_string))
@@ -110,7 +123,7 @@ class TemplateGenerator:
         """
         template = self._generate_template_for_model_class(model.__class__)
 
-        for name in model.__fields_set__:
+        for name in model.model_fields_set:
             value = getattr(model, name)
             template[name] = self._generate_template_for_value(value)
 
@@ -130,19 +143,25 @@ class TemplateGenerator:
         """
         template: Dict[str, Any] = {}
 
-        for name, field in model_class.__fields__.items():
-            if self._is_model_class(field.outer_type_):
-                template[name] = self._generate_template_for_model_class(
-                    field.outer_type_
-                )
-            elif field.outer_type_ is Optional and self._is_model_class(
-                field.type_
-            ):
-                template[name] = self._generate_template_for_model_class(
-                    field.type_
-                )
-            else:
-                template[name] = field._type_display()
+        for name, field in model_class.model_fields.items():
+            annotation = field.annotation
+
+            if annotation is not None:
+                if self._is_model_class(annotation):
+                    template[name] = self._generate_template_for_model_class(
+                        annotation
+                    )
+
+                elif typing_utils.is_optional(
+                    annotation
+                ) and self._is_model_class(
+                    typing_utils.get_args(annotation)[0]
+                ):
+                    template[name] = self._generate_template_for_model_class(
+                        typing_utils.get_args(annotation)[0]
+                    )
+                else:
+                    template[name] = pydantic_repr.display_as_type(annotation)
 
         return template
 
@@ -188,12 +207,16 @@ class YAMLSerializationMixin(BaseModel):
 
         Args:
             sort_keys: Whether to sort the keys in the YAML representation.
-            **kwargs: Kwargs to pass to the pydantic json(...) method.
+            **kwargs: Kwargs to pass to the pydantic model_dump(...) method.
 
         Returns:
             YAML string representation.
         """
-        dict_ = json.loads(self.json(**kwargs, sort_keys=sort_keys))
+        dict_ = json.loads(
+            json.dumps(
+                self.model_dump(mode="json", **kwargs), sort_keys=sort_keys
+            )
+        )
         return yaml.dump(dict_, sort_keys=sort_keys)
 
     @classmethod
@@ -207,12 +230,12 @@ class YAMLSerializationMixin(BaseModel):
             The model instance.
         """
         dict_ = yaml_utils.read_yaml(path)
-        return cls.parse_obj(dict_)
+        return cls.model_validate(dict_)
 
 
 def validate_function_args(
     __func: Callable[..., Any],
-    __config: Dict[str, Any],
+    __config: Optional[ConfigDict],
     *args: Any,
     **kwargs: Any,
 ) -> Dict[str, Any]:
@@ -233,40 +256,225 @@ def validate_function_args(
     Returns:
         The validated arguments.
     """
-    parameter_prefix = "zenml__"
-
     signature = inspect.signature(__func)
-    parameters = [
-        param.replace(name=f"{parameter_prefix}{param.name}")
-        for param in signature.parameters.values()
-    ]
-    signature = signature.replace(parameters=parameters)
-
-    def f() -> None:
-        pass
-
-    # We create a dummy function with the original function signature, but
-    # add a prefix to all arguments to avoid potential clashes with pydantic
-    # BaseModel attributes
+
+    validated_args = ()
+    validated_kwargs = {}
+
+    def f(*args: Any, **kwargs: Dict[Any, Any]) -> None:
+        nonlocal validated_args
+        nonlocal validated_kwargs
+
+        validated_args = args
+        validated_kwargs = kwargs
+
+    # We create a dummy function with the original function signature to run
+    # pydantic validation without actually running the function code
     f.__signature__ = signature  # type: ignore[attr-defined]
-    f.__annotations__ = {
-        f"{parameter_prefix}{key}": annotation
-        for key, annotation in __func.__annotations__.items()
-    }
-
-    validation_func = ValidatedFunction(f, config=__config)
-
-    kwargs = {
-        f"{parameter_prefix}{key}": value for key, value in kwargs.items()
-    }
-    model = validation_func.init_model_instance(*args, **kwargs)
-
-    validated_args = {
-        k[len(parameter_prefix) :]: v
-        for k, v in model._iter()
-        if k in model.__fields_set__
-        or model.__fields__[k].default_factory
-        or model.__fields__[k].default
-    }
-
-    return validated_args
+    f.__annotations__ = __func.__annotations__
+
+    validated_function = validate_call(config=__config, validate_return=False)(
+        f
+    )
+
+    # This raises a pydantic.ValidatonError in case the arguments are not valid
+    validated_function(*args, **kwargs)
+
+    return signature.bind(*validated_args, **validated_kwargs).arguments
+
+
+def model_validator_data_handler(
+    raw_data: Any,
+    base_class: Type[BaseModel],
+    validation_info: ValidationInfo,
+) -> Dict[str, Any]:
+    """Utility function to parse raw input data of varying types to a dict.
+
+    With the change to pydantic v2, validators which operate with "before"
+    (or previously known as the "pre" parameter) are getting "Any" types of raw
+    input instead of a "Dict[str, Any]" as before. Depending on the use-case,
+    this can create conflicts after the migration and this function will be
+    used as a helper function to handle different types of raw input data.
+
+    A code snippet to showcase how the behaviour changes. The "before" validator
+    prints the type of the input:
+
+        class Base(BaseModel):
+            a: int = 3
+
+        class MyClass(Base):
+            @model_validator(mode="before")
+            @classmethod
+            def before_validator(cls, data: Any) -> Any:
+                print(type(data))
+                return {}
+
+        one = MyClass() # prints "<class 'dict'>"
+        MyClass.model_validate(one)  # prints NOTHING, it is already validated
+        MyClass.model_validate("asdf")  # prints "<class 'str'>", fails without the modified return.
+        MyClass.model_validate(RandomClass())  # prints "<class 'RandomClass'>", fails without the modified return.
+        MyClass.model_validate(Base())  # prints "<class 'Base'>", fails without the modified return.
+        MyClass.model_validate_json(json.dumps("aria"))  # prints "<class 'str'>", fails without the modified return.
+        MyClass.model_validate_json(json.dumps([1]))  # prints "<class 'list'>", fails without the modified return.
+        MyClass.model_validate_json(one.model_dump_json())  # prints "<class 'dict'>"
+
+    Args:
+        raw_data: The raw data passed to the validator, can be "Any" type.
+        base_class: The class that the validator belongs to
+        validation_info: Extra information about the validation process.
+
+    Raises:
+        TypeError: if the type of the data is not processable.
+        ValueError: in case of an unknown validation mode.
+
+    Returns:
+        A dictionary which will be passed to the eventual validator of pydantic.
+    """
+    if validation_info.mode == "python":
+        # This is mode is only active if people validate objects using pythonic
+        # raw data such as MyClass(...) or MyClass.model_validate()
+
+        if isinstance(raw_data, dict):
+            # In most cases, this is the behaviour as the raw input is a dict
+            return raw_data
+
+        elif isinstance(raw_data, base_class):
+            # In some cases, we pass the same object type to the validation
+            # in such cases, it is critical we keep the original structure of
+            # fields that are already set.
+            return dict(raw_data)
+
+        elif issubclass(base_class, raw_data.__class__):
+            # There are a few occurrences where the annotation of the field is
+            # denoted by a subclass, and we use the instance of its super class
+            # as the raw input. In such cases we will use the same approach as
+            # before, while raising a debug message.
+            logger.debug(
+                f"During the validation of a `{base_class}` object, an instance"
+                f"of `{raw_data.__class__}` (super class of `{base_class}`) "
+                f"has been passed as raw input. This might lead to unexpected "
+                f"behaviour in case `{base_class}` have features which can not"
+                f"be extracted from an instance of a `{raw_data.__class__}`."
+            )
+            return dict(raw_data)
+
+        elif isinstance(raw_data, str):
+            # If the raw input is a raw string, we can try to use the `json`
+            # module to parse it. The resulting data needs to be a proper
+            # dict for us to pass it to the validation process.
+            try:
+                json_data = json.loads(raw_data)
+
+                if isinstance(json_data, dict):
+                    return json_data
+                else:
+                    raise TypeError("The resulting json data is not a dict!")
+
+            except (TypeError, JSONDecodeError) as e:
+                raise TypeError(
+                    "The raw json input string can not be converted to a "
+                    f"dict: {e}"
+                )
+        else:
+            raise TypeError(
+                "Unsupported type of raw input data for the `python` validation"
+                "mode of the pydantic class. Please consider changing the way "
+                f"you are creating using the `{base_class}` or instead use"
+                f"`{base_class}.model_validate_json()`."
+            )
+
+    elif validation_info.mode == "json":
+        # This is mode is only active if people validate objects using json
+        # input data such as MyClass.model_validate_json()
+        if isinstance(raw_data, dict):
+            return raw_data
+        else:
+            raise TypeError(
+                f"The resulting JSON data {raw_data} is not a dict, therefore"
+                f"can not be used by the validation process."
+            )
+    else:
+        # Unknown validation mode
+        raise ValueError(f"Unknown validation mode. {validation_info.mode}")
+
+
+def before_validator_handler(
+    method: Callable[..., Any],
+) -> Callable[[Any, Any, Any], Any]:
+    """Decorator to handle the raw input data for pydantic model validators.
+
+    Args:
+        method: the class method with the actual validation logic.
+
+    Returns:
+        the validator method
+    """
+
+    def before_validator(
+        cls: Type[BaseModel], data: Any, validation_info: ValidationInfo
+    ) -> Any:
+        """Wrapper method to handle the raw data.
+
+        Args:
+            cls: the class handler
+            data: the raw input data
+            validation_info: the context of the validation.
+
+        Returns:
+            the validated data
+        """
+        data = model_validator_data_handler(
+            raw_data=data, base_class=cls, validation_info=validation_info
+        )
+        return method(cls=cls, data=data)
+
+    return before_validator
+
+
+def has_validators(
+    pydantic_class: Type[BaseModel],
+    field_name: Optional[str] = None,
+) -> bool:
+    """Function to check if a Pydantic model or a pydantic field has validators.
+
+    Args:
+        pydantic_class: The class defining the pydantic model.
+        field_name: Optional, field info. If specified, this function will focus
+            on a singular field within the class. If not specified, it will
+            check model validators.
+
+    Returns:
+        Whether the specified field or class has a validator or not.
+    """
+    # If field is not specified check model validators
+    if field_name is None:
+        if pydantic_class.__pydantic_decorators__.model_validators:
+            return True
+
+    # Else, check field validators
+    else:
+        # 1. Field validators can be defined through @field_validator decorators
+        f_validators = pydantic_class.__pydantic_decorators__.field_validators
+
+        for name, f_v in f_validators.items():
+            if field_name in f_v.info.fields:
+                return True
+
+        # 2. Field validators can be defined through the Annotation[.....]
+        field_info = pydantic_class.model_fields[field_name]
+        if metadata := field_info.metadata:
+            if any(
+                isinstance(
+                    m,
+                    (
+                        AfterValidator,
+                        BeforeValidator,
+                        PlainValidator,
+                        WrapValidator,
+                    ),
+                )
+                for m in metadata
+            ):
+                return True
+
+    return False

zenml/utils/secret_utils.py

@@ -16,16 +16,29 @@
 import re
 from typing import TYPE_CHECKING, Any, NamedTuple
 
-from pydantic import Field
+from pydantic import Field, PlainSerializer, SecretStr
+from typing_extensions import Annotated
+
+from zenml.logger import get_logger
 
 if TYPE_CHECKING:
-    from pydantic.fields import ModelField
+    from pydantic.fields import FieldInfo
 
 _secret_reference_expression = re.compile(r"\{\{\s*\S+?\.\S+\s*\}\}")
 
 PYDANTIC_SENSITIVE_FIELD_MARKER = "sensitive"
 PYDANTIC_CLEAR_TEXT_FIELD_MARKER = "prevent_secret_reference"
 
+PlainSerializedSecretStr = Annotated[
+    SecretStr,
+    PlainSerializer(
+        lambda v: v.get_secret_value() if v is not None else None,
+        when_used="json",
+    ),
+]
+
+logger = get_logger(__name__)
+
 
 def is_secret_reference(value: Any) -> bool:
     """Checks whether any value is a secret reference.
@@ -87,8 +100,9 @@ def SecretField(*args: Any, **kwargs: Any) -> Any:
     Returns:
         Pydantic field info.
     """
-    kwargs[PYDANTIC_SENSITIVE_FIELD_MARKER] = True
-    return Field(*args, **kwargs)  # type: ignore[pydantic-field]
+    json_schema_extra = kwargs.get("json_schema_extra", {})
+    json_schema_extra.update({PYDANTIC_SENSITIVE_FIELD_MARKER: True})
+    return Field(json_schema_extra=json_schema_extra, *args, **kwargs)  # type: ignore[pydantic-field]
 
 
 def ClearTextField(*args: Any, **kwargs: Any) -> Any:
@@ -103,11 +117,12 @@ def ClearTextField(*args: Any, **kwargs: Any) -> Any:
     Returns:
         Pydantic field info.
     """
-    kwargs[PYDANTIC_CLEAR_TEXT_FIELD_MARKER] = True
-    return Field(*args, **kwargs)  # type: ignore[pydantic-field]
+    json_schema_extra = kwargs.get("json_schema_extra", {})
+    json_schema_extra.update({PYDANTIC_CLEAR_TEXT_FIELD_MARKER: True})
+    return Field(json_schema_extra=json_schema_extra, *args, **kwargs)  # type: ignore[pydantic-field]
 
 
-def is_secret_field(field: "ModelField") -> bool:
+def is_secret_field(field: "FieldInfo") -> bool:
     """Returns whether a pydantic field contains sensitive information or not.
 
     Args:
@@ -116,10 +131,29 @@ def is_secret_field(field: "ModelField") -> bool:
     Returns:
         `True` if the field contains sensitive information, `False` otherwise.
     """
-    return field.field_info.extra.get(PYDANTIC_SENSITIVE_FIELD_MARKER, False)  # type: ignore[no-any-return]
-
-
-def is_clear_text_field(field: "ModelField") -> bool:
+    if field.json_schema_extra is not None:
+        if isinstance(field.json_schema_extra, dict):
+            if marker := field.json_schema_extra.get(
+                PYDANTIC_SENSITIVE_FIELD_MARKER
+            ):
+                assert isinstance(marker, bool), (
+                    f"The parameter `{PYDANTIC_SENSITIVE_FIELD_MARKER}` in the "
+                    f"field definition can only be a boolean value."
+                )
+                return marker
+
+        else:
+            logger.warning(
+                f"The 'json_schema_extra' of the field '{field.title}' is "
+                "not defined as a dict. This might lead to unexpected "
+                "behaviour as we are checking it is a secret text field. "
+                "Returning 'False' as default..."
+            )
+
+    return False
+
+
+def is_clear_text_field(field: "FieldInfo") -> bool:
     """Returns whether a pydantic field prevents secret references or not.
 
     Args:
@@ -128,4 +162,23 @@ def is_clear_text_field(field: "ModelField") -> bool:
     Returns:
         `True` if the field prevents secret references, `False` otherwise.
     """
-    return field.field_info.extra.get(PYDANTIC_CLEAR_TEXT_FIELD_MARKER, False)  # type: ignore[no-any-return]
+    if field.json_schema_extra is not None:
+        if isinstance(field.json_schema_extra, dict):
+            if marker := field.json_schema_extra.get(
+                PYDANTIC_CLEAR_TEXT_FIELD_MARKER
+            ):
+                assert isinstance(marker, bool), (
+                    f"The parameter `{PYDANTIC_CLEAR_TEXT_FIELD_MARKER}` in the "
+                    f"field definition can only be a boolean value."
+                )
+                return marker
+
+        else:
+            logger.warning(
+                f"The 'json_schema_extra' of the field '{field.title}' is "
+                "not defined as a dict. This might lead to unexpected "
+                "behaviour as we are checking it is a clear text field. "
+                "Returning 'False' as default..."
+            )
+
+    return False

zenml/utils/source_utils.py

@@ -83,11 +83,11 @@ def load(source: Union[Source, str]) -> Any:
 
     import_root = None
     if source.type == SourceType.CODE_REPOSITORY:
-        source = CodeRepositorySource.parse_obj(source)
+        source = CodeRepositorySource.model_validate(dict(source))
         _warn_about_potential_source_loading_issues(source=source)
         import_root = get_source_root()
     elif source.type == SourceType.DISTRIBUTION_PACKAGE:
-        source = DistributionPackageSource.parse_obj(source)
+        source = DistributionPackageSource.model_validate(dict(source))
         if source.version:
             current_package_version = _get_package_version(
                 package_name=source.package_name

zenml/utils/typed_model.py

@@ -17,7 +17,9 @@ import json
 from typing import Any, Dict, Tuple, Type, cast
 
 from pydantic import BaseModel, Field
-from pydantic.main import ModelMetaclass
+
+# TODO: Investigate if we can solve this import a different way.
+from pydantic._internal._model_construction import ModelMetaclass
 from typing_extensions import Literal
 
 from zenml.utils import source_utils
@@ -85,7 +87,7 @@ class BaseTypedModel(BaseModel, metaclass=BaseTypedModelMeta):
 
     matrix = TheMatrix(choice=RedPill())
     d = matrix.dict()
-    new_matrix = TheMatrix.parse_obj(d)
+    new_matrix = TheMatrix.model_validate(d)
     assert isinstance(new_matrix.choice, RedPill)
     ```
 
@@ -126,7 +128,7 @@ class BaseTypedModel(BaseModel, metaclass=BaseTypedModelMeta):
                 f"Class `{cls}` is not a ZenML BaseTypedModel subclass."
             )
 
-        return cls.parse_obj(model_dict)
+        return cls.model_validate(model_dict)
 
     @classmethod
     def from_json(