From 48725944f4bf4705b6f35438f6bc54fcfff51f00 Mon Sep 17 00:00:00 2001 From: Michael Williamson Date: Mon, 3 Apr 2017 11:46:40 +0100 Subject: Add transforms.get_descendants() --- mammoth/transforms.py | 19 +++++++++++++++++++ tests/transforms_tests.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) create mode 100644 tests/transforms_tests.py 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)) -- cgit v1.2.3