diff options
| -rw-r--r-- | mammoth/transforms.py | 7 | ||||
| -rw-r--r-- | tests/transforms_tests.py | 12 |
2 files changed, 18 insertions, 1 deletions
diff --git a/mammoth/transforms.py b/mammoth/transforms.py index 88db8a1..21bcbf6 100644 --- a/mammoth/transforms.py +++ b/mammoth/transforms.py @@ -22,6 +22,13 @@ def _each_element(transform_element): return transform_element_and_children +def get_descendants_of_type(element, element_type): + return list(filter( + lambda descendant: isinstance(descendant, element_type), + get_descendants(element), + )) + + def get_descendants(element): descendants = [] diff --git a/tests/transforms_tests.py b/tests/transforms_tests.py index 128b9cb..e4f5360 100644 --- a/tests/transforms_tests.py +++ b/tests/transforms_tests.py @@ -1,7 +1,7 @@ from nose.tools import assert_equal, istest from mammoth import documents -from mammoth.transforms import get_descendants +from mammoth.transforms import get_descendants, get_descendants_of_type @istest @@ -26,3 +26,13 @@ class GetDescendantsTests(object): child = documents.run(children=[grandchild]) element = documents.paragraph(children=[child]) assert_equal([grandchild, child], get_descendants(element)) + + +@istest +class GetDescendantsOfTypeTests(object): + @istest + def filters_descendants_to_type(self): + tab = documents.tab() + run = documents.run(children=[]) + element = documents.paragraph(children=[tab, run]) + assert_equal([run], get_descendants_of_type(element, documents.Run)) |