Last active
July 18, 2024 01:10
-
-
Save shashankrnr32/c4e68520b46b2a3fea6f933cac4159ca to your computer and use it in GitHub Desktop.
Pydantic Model with Builder pattern
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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) |
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
Both of the below configurations should work fine