Skip to content

Instantly share code, notes, and snippets.

@shashankrnr32
Last active July 18, 2024 01:10
Show Gist options
  • Save shashankrnr32/c4e68520b46b2a3fea6f933cac4159ca to your computer and use it in GitHub Desktop.
Save shashankrnr32/c4e68520b46b2a3fea6f933cac4159ca to your computer and use it in GitHub Desktop.
Pydantic Model with Builder pattern
"""
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)
@shashankrnr32
Copy link
Author

shashankrnr32 commented Mar 10, 2021

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())

@shashankrnr32
Copy link
Author

shashankrnr32 commented Apr 6, 2021

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')

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment