Rework iterator on form

pydantic_form/iterators.py

@@ -1,22 +1,23 @@
-from wtforms import Form, FormField
+from wtforms import Form, FormField, FieldList
 from pydantic.fields import Shape
 
 from .utils import *
 
 
-def iter_field(parent_class, leafs_only=True, path=()):
-    field = getattr(parent_class, path[-1]) if path else parent_class
-
+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(field_class, leafs_only, path + (subfield_name,))
+            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
+        yield path, field, parent
 
 
 def iter_form_class(form_class, leafs_only=True, path=()):
@@ -27,23 +28,21 @@ def iter_form_class(form_class, leafs_only=True, path=()):
 
     try:
         for subfield_name, subfield in field._unbound_fields:
-            yield from iter_field(field, leafs_only, path + (subfield_name,))
+            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, path=()):
-    field = getattr(form, path[-1]) if path else form
-    if isinstance(field, Form) or isinstance(field, FormField):
-        if path and not leafs_only:
-            yield path, field
+        yield from iter_form_class(field, leafs_only, path)
 
-        for f in field._fields:
-            yield from iter_form(field, leafs_only, path + (f,))
-    else:
-        yield path, field
+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_

tests/test_iterators.py

@@ -51,8 +51,8 @@ def test_iterator_schema(scenario):
 def test_iterator_form(scenario):
     keys = scenario.keys
 
-    assert [k for k, _ in iter_form_class(scenario.form)] == keys.class_
+    assert [k for k, _, _ in iter_form_class(scenario.form)] == keys.class_
 
     form = scenario.form()
-    assert [k for k, _ in iter_form_class(scenario.form)] == keys.class_
+    assert [k for k, _, _ in iter_form_class(scenario.form)] == keys.class_
     assert [k for k, _ in iter_form(form)] == keys.instance