11

What I Have

Dictionary:

user_dict = {'user': {'field1': 'value1',
                      'field2': 'value2'},
             'admin':{'field1': 'value3',
                      'field2': 'value4'}}

Pydantic Model:

class User(BaseModel):
    account_type: Optional[str] = 'user'
    field1: Optional[str] = ''
    field1: Optional[str] = ''

    class Config:
        validate_assignment = True

    @validator("account_type", pre=True, always=True)
        def _set_account_type(cls, account_type: str):
        return account_type or "user"

Desired Outcome

I would like to lookup the values of field1 and field2 from the user_dict based upon the account_type. So I should be able to do something like this:

user = User()
print(user.field1)
# > 'value1'

print(user.field2)
# > 'value2'

user2 = User(field1=None, field2=None)
print(user2.field1)
# > 'value1'

print(user2.field2)
# > 'value2'

user3 = User(account_type="admin")
print(user3.field1)
# > 'value3'

print(user3.field2)
# > 'value4'

What I have tried

class User(BaseModel):
    account_type: Optional[str] = 'user'
    field1: Optional[str] = user_dict['account_type']['field1']
    field1: Optional[str] = user_dict['account_type']['field2']

    class Config:
        validate_assignment = True

    @validator("account_type", pre=True, always=True)
        def _set_account_type(cls, account_type: str):
        return account_type or "user"

user = User()
print(user.field1)
# > 'value1'

user = User(field1=None)
print(user.field1)
# > None

Using Validators for field1 and field2:

class User(BaseModel):
    account_type: Optional[str] = 'user'
    field1: Optional[str]
    field1: Optional[str]

    class Config:
        validate_assignment = True

    @validator("account_type", pre=True, always=True)
        def _set_account_type(cls, account_type: str):
        return account_type or "user"
    
    @validator("field1", pre=True, always=True)
        def _set_plan_credits(cls, value):
            if value is None:
                value = 'value1'
        return user_dict['account_type']['field1'] or value

    @validator("field2", pre=True, always=True)
        def _set_rollover_credits(cls, value):
            if value is None:
                value = 'value2'
        return user_dict['account_type']['field2'] or value

user = User()
# > AttributeError: type object 'Account' has no attribute 'account_type'

Any ideas on what I can do to get the desired outcome?

Manas Sambare
  • 1,158
  • 2
  • 11
  • 22

1 Answers1

11

I solved it by using the root_validator decorator as follows:

Solution:

@root_validator(pre=False)
def _set_fields(cls, values: dict) -> dict:
    """This is a validator that sets the field values based on the
    the user's account type.

    Args:
        values (dict): Stores the attributes of the User object.

    Returns:
        dict: The attributes of the user object with the user's fields.
    """
    values["field1"] = user_dict[values["account_type"]]["field1"]
    values["field2"] = user_dict[values["account_type"]]["field2"]
    return values
Manas Sambare
  • 1,158
  • 2
  • 11
  • 22
  • Can we use this root_validator to return some extra data in the response model as we do in Django serializer to get some extra data in the response – mdhv_kothari Jan 14 '23 at 12:55