Python pydantic.validator() Examples

def test_assert_raises_validation_error():
    test_name = test_assert_raises_validation_error.__name__

    class Model(BaseModel):
        a: str

        @validator('a')
        def check_a(cls, v):
            assert v == 'a', 'invalid a'
            return v

    Model(a='a')
    expected_errors = [{'loc': ('a',), 'msg': 'invalid a', 'type': 'assertion_error'}]

    try:
        Model(a='snap')
    except ValidationError as exc:
        actual_errors = exc.errors()
        if actual_errors != expected_errors:
            raise RuntimeError(f'{test_name}:\nActual errors: {actual_errors}\nExpected errors: {expected_errors}')
    else:
        raise RuntimeError(f'{test_name}: ValidationError was not raised') 

def test_validate_pre():
    @dataclass
    class MyDataclass:
        a: List[int]

        @validator('a', pre=True)
        def check_a1(cls, v):
            v.append('123')
            return v

        @validator('a')
        def check_a2(cls, v):
            v.append(456)
            return v

    assert MyDataclass(a=[1, 2]).a == [1, 2, 123, 456] 

def test_error_on_optional():
    class Foobar(BaseModel):
        foo: Optional[str] = None

        @validator('foo', always=True, pre=True)
        def check_foo(cls, v):
            raise ValueError('custom error')

    with pytest.raises(ValidationError) as exc_info:
        Foobar(foo='x')
    assert exc_info.value.errors() == [{'loc': ('foo',), 'msg': 'custom error', 'type': 'value_error'}]
    assert repr(exc_info.value.raw_errors[0]) == "ErrorWrapper(exc=ValueError('custom error'), loc=('foo',))"

    with pytest.raises(ValidationError) as exc_info:
        Foobar(foo=None)
    assert exc_info.value.errors() == [{'loc': ('foo',), 'msg': 'custom error', 'type': 'value_error'}] 

def test_validate_assignment_value_change():
    class Config:
        validate_assignment = True

    @pydantic.dataclasses.dataclass(config=Config, frozen=False)
    class MyDataclass:
        a: int

        @validator('a')
        def double_a(cls, v):
            return v * 2

    d = MyDataclass(2)
    assert d.a == 4

    d.a = 3
    assert d.a == 6 

def test_inheritance_replace():
    @dataclass
    class Parent:
        a: int

        @validator('a')
        def add_to_a(cls, v):
            return v + 1

    @dataclass
    class Child(Parent):
        @validator('a')
        def add_to_a(cls, v):
            return v + 5

    assert Child(a=0).a == 5 

def test_field_with_validator():
    class Model(BaseModel):
        something: Optional[int] = None

        @validator('something')
        def check_field(cls, v, *, values, config, field):
            return v

    assert Model.schema() == {
        'title': 'Model',
        'type': 'object',
        'properties': {'something': {'type': 'integer', 'title': 'Something'}},
    } 

def reload_config():
    global config

    class FullConfig(BaseModel):
        DataRoot: str = '/root/coolq/data/image'
        LogRoot: str = '/var/log/umr'
        CommandPrefix: str = '!!'
        Extensions: Optional[List[str]]
        BotAdmin: Optional[Dict[str, List[Union[int, str]]]]
        LogLevel: Optional[Dict[str, LogLevel]]

        ForwardList: ForwardList
        Driver: Optional[Dict[str, construct_union(driver_config, BaseDriverConfig)]]

        ExtensionConfig: Optional[Dict[str, construct_union(extension_config, BaseExtensionConfig)]]

        @validator('Extensions', pre=True, always=True)
        def generate_empty_list_if_none(cls, v):
            return v or []

        @validator('Driver', pre=True, always=True)
        def generate_empty_dict_if_none(cls, v):
            return v or {}

        @validator('ExtensionConfig', pre=True, always=True)
        def generate_empty_dict_if_none2(cls, v):
            return v or {}

        @validator('BotAdmin', pre=True, always=True)
        def generate_empty_dict_if_none3(cls, v):
            return v or {}

        @validator('LogLevel', pre=True, always=True)
        def generate_empty_dict_if_none4(cls, v):
            return v or {}

    config = FullConfig(**yaml.load(open(f'{home}/.umr/config.yaml'), yaml.FullLoader)) 

def test_inheritance_validators_all():
    class BarModel(BaseModel):
        @validator('*')
        def check_all(cls, v):
            return v * 2

    model = create_model('FooModel', a=(int, ...), b=(int, ...), __base__=BarModel)
    assert model(a=2, b=6).dict() == {'a': 4, 'b': 12} 

def test_inheritance_validators_always():
    class BarModel(BaseModel):
        @validator('a', check_fields=False, always=True)
        def check_a(cls, v):
            if 'foobar' not in v:
                raise ValueError('"foobar" not found in a')
            return v

    model = create_model('FooModel', a='cake', __base__=BarModel)
    with pytest.raises(ValidationError):
        model()
    assert model(a='this is foobar good').a == 'this is foobar good'
    with pytest.raises(ValidationError):
        model(a='something else') 

def test_inheritance_validators():
    class BarModel(BaseModel):
        @validator('a', check_fields=False)
        def check_a(cls, v):
            if 'foobar' not in v:
                raise ValueError('"foobar" not found in a')
            return v

    model = create_model('FooModel', a='cake', __base__=BarModel)
    assert model().a == 'cake'
    assert model(a='this is foobar good').a == 'this is foobar good'
    with pytest.raises(ValidationError):
        model(a='something else') 

def test_root_validator():
    root_val_values = []

    @dataclass
    class MyDataclass:
        a: int
        b: str

        @validator('b')
        def repeat_b(cls, v):
            return v * 2

        @root_validator
        def root_validator(cls, values):
            root_val_values.append(values)
            if 'snap' in values.get('b', ''):
                raise ValueError('foobar')
            return dict(values, b='changed')

    assert asdict(MyDataclass(a='123', b='bar')) == {'a': 123, 'b': 'changed'}

    with pytest.raises(ValidationError) as exc_info:
        MyDataclass(a=1, b='snap dragon')
    assert root_val_values == [{'a': 123, 'b': 'barbar'}, {'a': 1, 'b': 'snap dragonsnap dragon'}]

    assert exc_info.value.errors() == [{'loc': ('__root__',), 'msg': 'foobar', 'type': 'value_error'}] 

def test_validate_parent():
    @dataclass
    class Parent:
        a: str

        @validator('a')
        def change_a(cls, v):
            return v + ' changed'

    @dataclass
    class Child(Parent):
        pass

    assert Parent(a='this is foobar good').a == 'this is foobar good changed'
    assert Child(a='this is foobar good').a == 'this is foobar good changed' 

def test_classmethod():
    @dataclass
    class MyDataclass:
        a: str

        @validator('a')
        def check_a(cls, v):
            assert cls is MyDataclass and is_dataclass(MyDataclass)
            return v

    m = MyDataclass(a='this is foobar good')
    assert m.a == 'this is foobar good'
    m.check_a('x') 

def test_simple():
    @dataclass
    class MyDataclass:
        a: str

        @validator('a')
        def change_a(cls, v):
            return v + ' changed'

    assert MyDataclass(a='this is foobar good').a == 'this is foobar good changed' 

def test_value_validation():
    T = TypeVar('T')

    class Response(GenericModel, Generic[T]):
        data: T

        @validator('data', each_item=True)
        def validate_value_nonzero(cls, v):
            if v == 0:
                raise ValueError('value is zero')
            return v

        @root_validator()
        def validate_sum(cls, values):
            if sum(values.get('data', {}).values()) > 5:
                raise ValueError('sum too large')
            return values

    assert Response[Dict[int, int]](data={1: '4'}).dict() == {'data': {1: 4}}
    with pytest.raises(ValidationError) as exc_info:
        Response[Dict[int, int]](data={1: 'a'})
    assert exc_info.value.errors() == [
        {'loc': ('data', 1), 'msg': 'value is not a valid integer', 'type': 'type_error.integer'}
    ]

    with pytest.raises(ValidationError) as exc_info:
        Response[Dict[int, int]](data={1: 0})
    assert exc_info.value.errors() == [{'loc': ('data', 1), 'msg': 'value is zero', 'type': 'value_error'}]

    with pytest.raises(ValidationError) as exc_info:
        Response[Dict[int, int]](data={1: 3, 2: 6})
    assert exc_info.value.errors() == [{'loc': ('__root__',), 'msg': 'sum too large', 'type': 'value_error'}] 

def check_names_not_empty(cls, v):
        assert v != '', 'Empty strings are not allowed.'
        return v


# This will NOT raise a ValidationError because the validator was not called 

def test_generic():
    data_type = TypeVar('data_type')
    error_type = TypeVar('error_type')

    class Result(GenericModel, Generic[data_type, error_type]):
        data: Optional[List[data_type]]
        error: Optional[error_type]
        positive_number: int

        @validator('error', always=True)
        def validate_error(cls, v: Optional[error_type], values: Dict[str, Any]) -> Optional[error_type]:
            if values.get('data', None) is None and v is None:
                raise ValueError('Must provide data or error')
            if values.get('data', None) is not None and v is not None:
                raise ValueError('Must not provide both data and error')
            return v

        @validator('positive_number')
        def validate_positive_number(cls, v: int) -> int:
            if v < 0:
                raise ValueError
            return v

    class Error(BaseModel):
        message: str

    class Data(BaseModel):
        number: int
        text: str

    success1 = Result[Data, Error](data=[Data(number=1, text='a')], positive_number=1)
    assert success1.dict() == {'data': [{'number': 1, 'text': 'a'}], 'error': None, 'positive_number': 1}
    assert repr(success1) == "Result[Data, Error](data=[Data(number=1, text='a')], error=None, positive_number=1)"

    success2 = Result[Data, Error](error=Error(message='error'), positive_number=1)
    assert success2.dict() == {'data': None, 'error': {'message': 'error'}, 'positive_number': 1}
    assert repr(success2) == "Result[Data, Error](data=None, error=Error(message='error'), positive_number=1)"
    with pytest.raises(ValidationError) as exc_info:
        Result[Data, Error](error=Error(message='error'), positive_number=-1)
    assert exc_info.value.errors() == [{'loc': ('positive_number',), 'msg': '', 'type': 'value_error'}]

    with pytest.raises(ValidationError) as exc_info:
        Result[Data, Error](data=[Data(number=1, text='a')], error=Error(message='error'), positive_number=1)
    assert exc_info.value.errors() == [
        {'loc': ('error',), 'msg': 'Must not provide both data and error', 'type': 'value_error'}
    ]

    with pytest.raises(ValidationError) as exc_info:
        Result[Data, Error](data=[Data(number=1, text='a')], error=Error(message='error'), positive_number=1)
    assert exc_info.value.errors() == [
        {'loc': ('error',), 'msg': 'Must not provide both data and error', 'type': 'value_error'}
    ] 

def validate_structure_features(cls, v, values):
        if [StructureFeatures(value) for value in sorted((_.value for _ in v))] != v:
            raise ValueError(
                f"structure_features MUST be sorted alphabetically, given value: {v}"
            )
        # disorder
        for species in values.get("species", []):
            if len(species.chemical_symbols) > 1:
                if StructureFeatures.DISORDER not in v:
                    raise ValueError(
                        f"{StructureFeatures.DISORDER.value} MUST be present when any one entry in species has a chemical_symbols list greater than one element"
                    )
                break
        else:
            if StructureFeatures.DISORDER in v:
                raise ValueError(
                    f"{StructureFeatures.DISORDER.value} MUST NOT be present, since all species' chemical_symbols lists are equal to or less than one element"
                )
        # assemblies
        if values.get("assemblies", None) is not None:
            if StructureFeatures.ASSEMBLIES not in v:
                raise ValueError(
                    f"{StructureFeatures.ASSEMBLIES.value} MUST be present, since the property of the same name is present"
                )
        else:
            if StructureFeatures.ASSEMBLIES in v:
                raise ValueError(
                    f"{StructureFeatures.ASSEMBLIES.value} MUST NOT be present, since the property of the same name is not present"
                )
        # site_attachments
        for species in values.get("species", []):
            # There is no need to also test "nattached",
            # since a Species validator makes sure either both are present or both are None.
            if getattr(species, "attached", None) is not None:
                if StructureFeatures.SITE_ATTACHMENTS not in v:
                    raise ValueError(
                        f"{StructureFeatures.SITE_ATTACHMENTS.value} MUST be present when any one entry in species includes attached and nattached"
                    )
                break
        else:
            if StructureFeatures.SITE_ATTACHMENTS in v:
                raise ValueError(
                    f"{StructureFeatures.SITE_ATTACHMENTS.value} MUST NOT be present, since no species includes the attached and nattached fields"
                )
        # implicit_atoms
        species_names = [_.name for _ in values.get("species", [])]
        for name in species_names:
            if name not in values.get("species_at_sites", []):
                if StructureFeatures.IMPLICIT_ATOMS not in v:
                    raise ValueError(
                        f"{StructureFeatures.IMPLICIT_ATOMS.value} MUST be present when any one entry in species is not represented in species_at_sites"
                    )
                break
        else:
            if StructureFeatures.IMPLICIT_ATOMS in v:
                raise ValueError(
                    f"{StructureFeatures.IMPLICIT_ATOMS.value} MUST NOT be present, since all species are represented in species_at_sites"
                )
        return v