Pydantic. Часть 2
В конце предыдущей статьи мы остановились примерно на следующем
from pydantic import BaseModel, ValidationError, Field
from pydantic import validator, root_validator
from typing import Tuple
from enum import Enum
class Flavor(str, Enum):
chocolate = 'chocolate'
vanilla = 'vanilla'
strawberry = 'strawberry'
mint = 'mint'
coffeee = 'coffee'
peanut_butter = 'peanut butter'
class Topping(str, Enum):
sprinkles = 'sprinkles'
hot_fudge = 'hot fudge'
cookies = 'cookies'
brownie = 'brownie'
whipped_cream = 'whipped cream'
strawberries = 'strawberries'
class Container(str, Enum):
cup = 'cup'
cone = 'cone'
waffle_cone = 'waffle cone'
class IceCreamMix(BaseModel):
name: str
container: Container
flavor: Flavor
toppings: Tuple[Topping, ...]
scoops: int = Field(..., gt=0, lt=5)
@validator('toppings')
def check_toppings(cls, toppings):
if len(toppings) > 4:
raise ValueError('Too many toppings')
return toppings
@root_validator
def check_cone_toppings(cls, values):
container = values.get('container')
toppings = values.get('toppings')
if container == Container.cone or container == Container.waffle_cone:
if Topping.hot_fudge in toppings:
raise ValueError('Cones cannot have hot fudge')
return values
def main():
try:
ice_cream_mix = IceCreamMix(
name = "PB&J",
container = Container.waffle_cone,
flavor = Flavor.peanut_butter,
toppings = (Topping.strawberries, Topping.brownie,Topping.sprinkles),
scoops = 2
)
print(ice_cream_mix.json())
except ValidationError as e:
print(e.json())
if __name__ == '__main__':
main()
Можно получить json схему класса IceCreamMix
print(IceCreamMix.schema())
python PydanticDemo.py
{'title': 'IceCreamMix', 'type': 'object', 'properties': {'name': {'title': 'Name', 'type': 'string'}, 'container': {'$ref': '#/definitions/Container'}, 'flavor': {'$ref': '#/definitions/Flavor'}, 'toppings': {'type': 'array', 'items': {'$ref': '#/definitions/Topping'}}, 'scoops': {'title': 'Scoops', 'exclusiveMinimum': 0, 'exclusiveMaximum': 5, 'type': 'integer'}}, 'required': ['name', 'container', 'flavor', 'toppings', 'scoops'], 'definitions': {'Container': {'title': 'Container', 'description': 'An enumeration.', 'enum': ['cup', 'cone', 'waffle cone'], 'type': 'string'}, 'Flavor': {'title': 'Flavor', 'description': 'An enumeration.', 'enum': ['chocolate', 'vanilla', 'strawberry', 'mint', 'coffee', 'peanut butter'], 'type': 'string'}, 'Topping': {'title': 'Topping', 'description': 'An enumeration.', 'enum': ['sprinkles', 'hot fudge', 'cookies', 'brownie', 'whipped cream', 'strawberries'], 'type': 'string'}}}
