Fix recursion in transforms

2 files changed, 31 insertions, 2 deletions

diff --git a/mammoth/transforms.py b/mammoth/transforms.py
index d126578..649b4db 100644
--- a/mammoth/transforms.py
+++ b/mammoth/transforms.py
@@ -18,7 +18,7 @@ def element_of_type(element_type, transform):
 def _each_element(transform_element):
     def transform_element_and_children(element):
         if isinstance(element, documents.HasChildren):
-            children = list(map(transform_element, element.children))
+            children = list(map(transform_element_and_children, element.children))
             element = element.copy(children=children)
         
         return transform_element(element)
diff --git a/tests/transforms_tests.py b/tests/transforms_tests.py
index e4f5360..0f0a05d 100644
--- a/tests/transforms_tests.py
+++ b/tests/transforms_tests.py
@@ -1,7 +1,36 @@
+import cobble
 from nose.tools import assert_equal, istest
 
 from mammoth import documents
-from mammoth.transforms import get_descendants, get_descendants_of_type
+from mammoth.transforms import get_descendants, get_descendants_of_type, _each_element
+
+
+@istest
+class EachElementTests(object):
+    @istest
+    def all_descendants_are_transformed(self):
+        @cobble.data
+        class Count(documents.HasChildren):
+            count = cobble.field()
+        
+        root = Count(count=None, children=[
+            Count(count=None, children=[
+                Count(count=None, children=[]),
+            ]),
+        ])
+        
+        current_count = [0]
+        def set_count(node):
+            current_count[0] += 1
+            return node.copy(count=current_count[0])
+        
+        result = _each_element(set_count)(root)
+        
+        assert_equal(Count(count=3, children=[
+            Count(count=2, children=[
+                Count(count=1, children=[]),
+            ]),
+        ]), result)
 
 
 @istest