pydantic_form/pydantic_form/iterators.py

from wtforms import Form, FormField, FieldList
from pydantic.fields import Shape

from .utils import *


def iter_field(field, leafs_only=True, path=(), parent=None):
    if issubclass(field.field_class, FormField):
        field_class = field.kwargs['form_class']
        if path and not leafs_only:
            yield path, field_class

        for subfield_name, subfield in field_class._unbound_fields:
            yield from iter_field(subfield, leafs_only, path + (subfield_name,), parent)

    elif issubclass(field.field_class, FieldList):
        yield from iter_field(field.args[0], leafs_only, path, path)

    else:
        yield path, field, parent


def iter_form_class(form_class, leafs_only=True, path=()):
    field = getattr(form_class, path[-1]) if path else form_class

    if path and not leafs_only:
        yield path, field

    try:
        for subfield_name, subfield in field._unbound_fields:
            yield from iter_field(subfield, leafs_only, path + (subfield_name,))
    except TypeError as e:
        # Ensure that the _unbound_fields is populated (that happens on first
        # instantiation)
        form_class()
        yield from iter_form_class(field, leafs_only, path)

def iter_form(form, leafs_only=True):
    for key, field, iterate_on in iter_form_class(form.__class__):
        if iterate_on:
            value = rgetattr(form, iterate_on)
            for n, entry in enumerate(value.entries):
                yield key + (n,), entry
            continue
        yield key, field

def iter_schema_class_field(field, leafs_only=True, path=(), parent=None):
    type_ = field.type_
    if hasattr(type_, '__fields__'):
        if path and not leafs_only:
            yield path, field, None

        for key, sub_field in type_.__fields__.items():
            yield from iter_schema_class_field(sub_field, leafs_only, path + (key,), field)
    else:
        iterate_on = None
        if parent and parent.shape != Shape.SINGLETON:
            iterate_on = path[:-1]
        elif field.shape != Shape.SINGLETON:
            iterate_on = path

        yield path, field, iterate_on

def iter_schema_class(schema, leafs_only=True, path=()):
    for key, field in schema.__fields__.items():
        yield from iter_schema_class_field(field, leafs_only, path + (key,))

def iter_schema(schema, leafs_only=True):
    for key, field, iterate_on in iter_schema_class(schema.__class__, leafs_only):
        if iterate_on:
            sub_values = rgetattr(schema, iterate_on)
            if isinstance(sub_values, list):
                for n, sub_value in enumerate(sub_values):
                    yield key + (n,), sub_value
            else:
                raise NotImplementedError()
            continue
        yield key, rgetattr(schema, key)