code to walk models

[tlambert03]

from typing import Iterator, Type, Union, get_origin, get_args
from pydantic import BaseModel


def walk_model(model: Type[BaseModel]) -> Iterator[Type[BaseModel]]:
    for field_name, field in model.__fields__.items():
        if isinstance(field.type_, type) and issubclass(field.type_, BaseModel):
            yield from walk_model(field.type_)
            yield field.type_
        elif get_origin(field.type_) is Union:
            for union_type in get_args(field.type_):
                if isinstance(union_type, type) and issubclass(
                    union_type, BaseModel
                ):
                    yield from walk_model(union_type)
                    yield union_type