diff options
| -rw-r--r-- | mammoth/transforms.py | 2 | ||||
| -rw-r--r-- | tests/transforms_tests.py | 31 |
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 |