Pydantic - Turn attribute immutable when defined as non-None for the first time-CodePudding

I have the following model:

from pydantic import BaseModel

class User(BaseModel):
    user_id: Optional[int] = None
    name: str

The user_id may be initially unknown (None), but when it is set to a non-None value then it should be immutable afterwards.

Examples:

user1 = User(name="John")
user1.user_id = 1    # user_id is set to a non-None value for the first time
# from now on, user_id must be immutable
user1.user_id = 10   # should raise a ValidationError

user2 = User(user_id=None, name="John")
user2.user_id = 2    # user_id is set to a non-None value for the first time
# from now on, user_id must be immutable
user2.user_id = 20   # should raise a ValidationError
user2.user_id = None # should raise a ValidationError

user3 = User(user_id=None, name="John")
user3.user_id = None    # that's ok, the value is still None
user3.user_id = 3       # user_id is now set to a non-None value
# from now on, user_id must be immutable
user3.user_id = 30      # should raise a ValidationError

user4 = User(user_id=4, name="John")
# from now on, user_id must be immutable
user4.user_id = 40      # should raise a ValidationError

This solution requires the user_id to be initialized and thus would make it impossible to assign user1.user_id = 1.

CodePudding user response：

I agree with @chepner that you should probably re-think your algorithm. It seems strange that you even have the possibility of re-assignment, if you know that this is not supposed to happen. In addition, you could just simply check before assigning. But it is hard to tell, what may be a better course of action without additional context.

To the matter at hand.

Pydantic model validation is stateless by design, as far as I know. All validators are class methods and there is no built-in way to have assignment validation depend on the state of the model instance being validated.

That being said, you can always construct a workaround using standard Python "dunder" magic, without getting too much in the way of Pydantic-specifics. Attribute assignment is done via __setattr__, even in the case of Pydantic models. We can hook into that method minimally and do our check there.

To be as much in line with the "Pydantic way", we should not raise a ValidationError there immediately. Instead we can simply assign a unique object that is guaranteed to be picked up by a validator down the line, so that we get a "clean" validation error.

Here is how this can be implemented:

from typing import Any, ClassVar

from pydantic import BaseModel, validator


class User(BaseModel):
    _USER_ID_SENTINEL: ClassVar[object] = object()

    user_id: int | None = None
    name: str

    class Config:
        validate_assignment = True

    def __setattr__(self, key: str, value: Any) -> None:
        if key == "user_id" and self.user_id is not None:
            value = self.__class__._USER_ID_SENTINEL
        super().__setattr__(key, value)

    @validator("user_id", pre=True)
    def ensure_no_reassignment(cls, v: Any) -> Any:
        if v is cls._USER_ID_SENTINEL:
            raise ValueError("Re-assignment of user_id not allowed")
        return v

Here is a full working test case:

from unittest import TestCase, main
from typing import Any, ClassVar

from pydantic import BaseModel, ValidationError, validator


class User(BaseModel):
    ... # see above


class Test(TestCase):
    def test_default_delayed_assignment(self) -> None:
        user = User(name="John")
        user.user_id = 1
        with self.assertRaises(ValidationError):
            user.user_id = 10
        with self.assertRaises(ValidationError):
            user.user_id = None

    def test_explicit_none_delayed_assignment(self) -> None:
        user = User(user_id=None, name="John")
        user.user_id = 2
        with self.assertRaises(ValidationError):
            user.user_id = 20
        with self.assertRaises(ValidationError):
            user.user_id = None

    def test_delayed_assignment_none_first(self) -> None:
        user = User(user_id=None, name="John")
        user.user_id = None
        user.user_id = 3
        with self.assertRaises(ValidationError):
            user.user_id = 30
        with self.assertRaises(ValidationError):
            user.user_id = None

    def test_init_assignment(self) -> None:
        user = User(user_id=4, name="John")
        with self.assertRaises(ValidationError):
            user.user_id = 40

    def test_normal_functionality(self) -> None:
        with self.assertRaises(ValidationError):
            User(**{"name": object()})
        with self.assertRaises(ValidationError):
            User.parse_obj({"name": "John", "user_id": object()})
        user = User(name="John")
        user.name = "Alice"
        self.assertDictEqual({"user_id": None, "name": "Alice"}, user.dict())


if __name__ == '__main__':
    main()

The pre=True in the @validator is necessary to allow us to give a useful error message. We could omit the custom validator entirely. Then the regular int validator will pick the object up and raise an error, but that error will be misleading, indicating a wrong type.

The validate_assignment = True setting in the Config is obviously necessary for a validator to even be called after initialization.

Also, if you wanted to have even more information in the error message, you could (for example) use a special class instead of object for the sentinel and substitute an instance of that class in __setattr__ containing the pre-existing value of user_id and the value trying to be assigned. Then you could pick that up in the validator method and include that information in the error message. I don't know if that is useful though.

Note that (as always) there is no real immutability because in Python there is always a way to mutate any attribute of any object. However, this solution is as "real" as the Pydantic concept of immutability.