pretty validation error

pretty.py

import re
import json

from pydantic import BaseModel, ValidationError, validator, conint

from typing import List, Literal

class Person(BaseModel):
    age: conint(ge=0, le=120)
    name: str

class Student(Person):
    subject: Literal['Math', 'English', 'History']

class School(BaseModel):
    name: str
    students: List[Student]

    @validator('name')
    def validate_name(cls, value: str):
        # Some stupid test
        if value.title() != value:
            raise ValueError('Name must be in title case')
        return value


def prettify_validation_error(error: ValidationError, include_extra_fields=False) -> str:
    return f'Found {len(error.errors())} validation errors for class {error.model.__name__}\n' \
           f'{json.dumps(format_error_dict(error.errors(), include_extra_fields), indent=2)}'


def format_error_dict(errors: List[dict], include_extra_fields):
    d = {'errors': []}
    for error in errors:
        error_dict = {
            'field': _format_location(error['loc']),
            'message': error['msg']
        }
        if include_extra_fields:
            error_dict['extra'] = {key: value for key, value in error.items() if key not in {'loc', 'msg'}}
        d['errors'].append(error_dict)

    return d


def _format_location(loc: tuple):
    """
    Formats pydantic's `loc` param to be more readable
    e.g (company, 1, size, min) -> 'company[1].size.min'
    """
    location = '.'.join([str(i) for i in loc])
    return re.sub(r'\.(\d+)', r'[\1]', location)  # replace xxx.1.xxx to xxx[1].xxx to indicate index

def main():
    data = {
        'name': 'school of fun',
        'students': [
            {
                'subject': 'Math',
                'age': 22,
                'name': 'Ron'
            },
            {
                'subject': 'Design',
                'age': 500,
                'name': 'John'
            },
            {
                'subject': 'Math',
                'age': 30,
                'name': 'Mark'
            },
        ]
    }

    # To get pretty output in stack trace
    #ValidationError.__str__ = lambda self: prettify_validation_error(self, True)
    #School.parse_obj(data)

    # As response to a client
    try:
        School.parse_obj(data)
    except ValidationError as e:
        print(prettify_validation_error(e, include_extra_fields=False))

if __name__ == '__main__':
    main()