Add transforms.get_descendants()

2 files changed, 47 insertions, 0 deletions

diff --git a/mammoth/transforms.py b/mammoth/transforms.py
index 6eb6d92..88db8a1 100644
--- a/mammoth/transforms.py
+++ b/mammoth/transforms.py
@@ -20,3 +20,22 @@ def _each_element(transform_element):
         return transform_element(element)
     
     return transform_element_and_children
+
+
+def get_descendants(element):
+    descendants = []
+
+    def visit(element):
+        descendants.append(element)
+
+    _visit_descendants(element, visit)
+
+    return descendants
+
+
+def _visit_descendants(element, visit):
+    if isinstance(element, documents.HasChildren):
+        for child in element.children:
+            _visit_descendants(child, visit)
+            visit(child)
+
diff --git a/tests/transforms_tests.py b/tests/transforms_tests.py
new file mode 100644
index 0000000..128b9cb
--- /dev/null
+++ b/tests/transforms_tests.py
@@ -0,0 +1,28 @@
+from nose.tools import assert_equal, istest
+
+from mammoth import documents
+from mammoth.transforms import get_descendants
+
+
+@istest
+class GetDescendantsTests(object):
+    @istest
+    def returns_nothing_if_element_type_has_no_children(self):
+        assert_equal([], get_descendants(documents.tab()))
+    
+    @istest
+    def returns_nothing_if_element_has_empty_children(self):
+        assert_equal([], get_descendants(documents.paragraph(children=[])))
+    
+    @istest
+    def includes_children(self):
+        children = [documents.text("child 1"), documents.text("child 2")]
+        element = documents.paragraph(children=children)
+        assert_equal(children, get_descendants(element))
+    
+    @istest
+    def includes_indirect_descendants(self):
+        grandchild = documents.text("grandchild")
+        child = documents.run(children=[grandchild])
+        element = documents.paragraph(children=[child])
+        assert_equal([grandchild, child], get_descendants(element))