Pydantic Model with Builder pattern

model.py

"""
Github Issue Thread: https://github.com/samuelcolvin/pydantic/issues/2152#issuecomment-786713976

License : MIT

Copyright 2021 Shashank Sharma

Permission is hereby granted, free of charge, to any person obtaining a copy of this software 
and associated documentation files (the "Software"), to deal in the Software without restriction, 
including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, 
and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, 
subject to the following conditions:

The above copyright notice and this permission notice shall be included in 
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, 
INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 
WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR 
IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

from functools import partial
from types import SimpleNamespace
from typing import Type

from pydantic import BaseModel, Extra, ValidationError


class ModelBuilder(object):
    """

    The Pydantic Model Builder class. Initialize the Builder with
    a custom Pydantic Model to get a Builder object for the particular model.

    You can set a field in the model in any of 2 ways below.

    1. builder.customField = customValue
    2. builder = builder.customField(customValue)

    The 2nd way above allows to chain the method calls.

    Examples:
        ```python
        model = builder.customField1(customValue1).customField2(customValue2).build()
        ```

    You can also turn off the field value validation by passing `validation=False`
    to the constructor.
    """

    __slots__ = ("__model__", "__values__", "__builder_config__")

    def __init__(self, model: Type[BaseModel], **kwargs) -> None:
        """
        Construct a Builder object for a Pydantic Model class

        Args:
            model: The Pydantic Model Class
            **kwargs: Keyword Arguments

        Keyword Args:
            validation (bool):
                Turn off field validation during setting. This ensures the model gets `construct`ed
                when `build()` method is called.
        """
        object.__setattr__(self, "__model__", model)
        object.__setattr__(self, "__values__", dict())

        # Builder configuration
        object.__setattr__(
            self,
            "__builder_config__",
            SimpleNamespace(validation=kwargs.pop("validation", True)),
        )

    def __getattr__(self, name: str):
        return partial(self.set_field, name=name)

    def __setattr__(self, key, value):
        self.set_field(value, key)

    def set_field(self, value: object, name: str, by_alias=False) -> "ModelBuilder":
        """
        Set the Pydantic field attribute. The value of the attribute can be another instance
        of `ModelBuilder`. The `build` method is called to update the fields

        Args:
            value:
                The value of the field. Value can be another builder.
                The `build` method will be called if the value is a subclass of ModelBuilder
            name: The name of the field
            by_alias: Setting field by alias

        Returns:
            The Builder object
        """
        # Check if the value is an instance of `ModelBuilder` call `build()` method to update value
        if issubclass(value.__class__, self.__class__) and hasattr(value, "build"):
            return self.set_field(value.build(), name)  # type: ignore

        if name not in self.__model__.__fields__:

            # If name is one of the aliases
            for field in self.__model__.__fields__.values():
                if field.alias == name:
                    return self.set_field(value, field.name, by_alias=True)

            if self.__model__.__config__.extra == Extra.forbid:
                raise KeyError(
                    f"'{name}' is not a valid field in '{self.__model__.__name__}'."
                )
            elif self.__model__.__config__.extra == Extra.ignore:
                return self
            elif self.__model__.__config__.extra == Extra.allow:
                self.__values__[name] = value
                return self

        known_field = self.__model__.__fields__[name]

        if not by_alias:
            if (
                known_field.alt_alias
                and not self.__model__.__config__.allow_population_by_field_name
                and name == known_field.name
            ):
                return self

        if self.__builder_config__.validation:
            _values_dict_excluding_alias = {
                k: v for k, v in self.__values__.items() if k != known_field.alias
            }
            value, error_ = known_field.validate(
                value,
                _values_dict_excluding_alias,
                loc=known_field.alias,
                cls=self.__model__.__class__,
            )
            if error_:
                raise ValidationError(
                    [
                        error_,
                    ],
                    self.__model__,
                )

        self.__values__[known_field.alias] = value
        return self

    def build(self) -> Type[BaseModel]:
        """
        Builds the Model. Calling this method returns a Model object with all the values
        passed to it.

        Returns:
            The Model object
        """
        if self.__builder_config__.validation:
            return self.__model__(**self.__values__)
        return self.__model__.construct(**self.__values__)


class BuildableBaseModel(BaseModel):
    """
    Adds capability to build a Pydantic model step by step instead of generating
    the entire object by passing the values to the constructor. Get the Model builder by calling the
    `Builder` method which in turn constructs an object of ModelBuilder type.
    """

    @classmethod
    def Builder(cls, **kwargs):
        return ModelBuilder(cls, **kwargs)

class MockModel(BuildableBaseModel):

    mock_str: str = Field(..., description="A mock pydantic field of string type", alias="bar")
    flag: bool = Field(..., description="A mock pydantic field of string type")

    class Config:
        allow_population_by_field_name = True
        extra = Extra.allow

Both of the below configurations should work fine

model = MockModel.Builder()
model.mock_str = "Hello"
model.flag = True
print(model.build())

model = MockModel.Builder().bar("Hello").flag(True)      # Even Alias works
print(model.build())

Other test code with nested models. Another ModelBuilder object can also be provided for which build method is called before updating the value.

from pydantic import Field
from typing import Optional


class FooModel(BuildableBaseModel):
    foo: str = Field(
        default="foovalue", description="A mock pydantic field of string type"
    )


class MockModel(BuildableBaseModel):

    mock_str: str = Field(
        ..., description="A mock pydantic field of string type", alias="bar"
    )
    flag: bool = Field(..., description="A mock pydantic field of string type")
    fooModel: Optional[FooModel] = Field(description="Mock Model attribute")

    class Config:
        allow_population_by_field_name = True
        extra = Extra.allow


model = (
    MockModel.Builder()
    .bar("Hello")
    .flag(True)
    .fooModel(FooModel.Builder().foo("testValue"))
)
print(model.build())
# mock_str='Hello' flag=True fooModel=FooModel(foo='testValue')