pydantic/tests/test_schema.py

from enum import Enum, IntEnum

import pytest

from pydantic import BaseModel, Schema, ValidationError


def test_key():
    class ApplePie(BaseModel):
        """
        This is a test.
        """
        a: float
        b: int = 10

    s = {
        'type': 'object',
        'title': 'ApplePie',
        'description': 'This is a test.',
        'properties': {
            'a': {
                'type': 'float',
                'required': True,
                'title': 'A',
            },
            'b': {
                'type': 'int',
                'required': False,
                'title': 'B',
                'default': 10,
            },
        },
    }
    assert True not in ApplePie._schema_cache
    assert False not in ApplePie._schema_cache
    assert ApplePie.schema() == s
    assert True in ApplePie._schema_cache
    assert False not in ApplePie._schema_cache
    assert ApplePie.schema() == s


def test_by_alias():
    class ApplePie(BaseModel):
        a: float
        b: int = 10

        class Config:
            title = 'Apple Pie'
            fields = {'a': 'Snap', 'b': 'Crackle'}

    s = {
        'type': 'object',
        'title': 'Apple Pie',
        'properties': {
            'Snap': {
                'type': 'float',
                'required': True,
                'title': 'Snap',
            },
            'Crackle': {
                'type': 'int',
                'required': False,
                'title': 'Crackle',
                'default': 10,
            },
        },
    }
    assert ApplePie.schema() == s
    assert ApplePie.schema() == s
    assert list(ApplePie.schema(by_alias=True)['properties'].keys()) == ['Snap', 'Crackle']
    assert list(ApplePie.schema(by_alias=False)['properties'].keys()) == ['a', 'b']


def test_sub_model():
    class Foo(BaseModel):
        """hello"""
        b: float

    class Bar(BaseModel):
        a: int
        b: Foo = None

    assert Bar.schema() == {
        'type': 'object',
        'title': 'Bar',
        'properties': {
            'a': {
                'type': 'int',
                'title': 'A',
                'required': True,
            },
            'b': {
                'type': 'object',
                'title': 'Foo',
                'properties': {
                    'b': {
                        'type': 'float',
                        'title': 'B',
                        'required': True,
                    },
                },
                'description': 'hello',
                'required': False,
            },
        },
    }


def test_schema_class():
    class Model(BaseModel):
        foo: int = Schema(4, title='Foo is Great')
        bar: str = Schema(..., description='this description of bar')

    with pytest.raises(ValidationError):
        Model()

    m = Model(bar=123)
    assert m.dict() == {'foo': 4, 'bar': '123'}

    assert Model.schema() == {
        'type': 'object',
        'title': 'Model',
        'properties': {
            'foo': {
                'type': 'int',
                'title': 'Foo is Great',
                'required': False,
                'default': 4,
            },
            'bar': {
                'type': 'str',
                'title': 'Bar',
                'required': True,
                'description': 'this description of bar',
            },
        },
    }


def test_schema_class_by_alias():
    class Model(BaseModel):
        foo: int = Schema(4, alias='foofoo')

    assert list(Model.schema()['properties'].keys()) == ['foofoo']
    assert list(Model.schema(by_alias=False)['properties'].keys()) == ['foo']


def test_choices():
    FooEnum = Enum('FooEnum', {'foo': 'f', 'bar': 'b'})
    BarEnum = IntEnum('BarEnum', {'foo': 1, 'bar': 2})

    class SpamEnum(str, Enum):
        foo = 'f'
        bar = 'b'

    class Model(BaseModel):
        foo: FooEnum
        bar: BarEnum
        spam: SpamEnum = Schema(None, choice_names={'f': 'Sausage'})

    assert Model.schema() == {
        'type': 'object',
        'title': 'Model',
        'properties': {
            'foo': {
                'type': 'enum',
                'title': 'Foo',
                'required': True,
                'choices': [
                    ('f', 'Foo'),
                    ('b', 'Bar'),
                ],
            },
            'bar': {
                'type': 'int',
                'title': 'Bar',
                'required': True,
                'choices': [
                    (1, 'Foo'),
                    (2, 'Bar'),
                ],
            },
            'spam': {
                'type': 'str',
                'title': 'Spam',
                'required': False,
                'choices': [
                    ('f', 'Sausage'),
                    ('b', 'Bar'),
                ],
            },
        },
    }