diff --git a/CHANGELOG b/CHANGELOG index 3832524..7d34015 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -1,5 +1,9 @@ v0.5 (unreleased): +- Added Wikicode.contains() to determine whether a Node or Wikicode object is + contained within another Wikicode object. +- Added Wikicode.get_ancestors() and Wikicode.get_parent() to find all + ancestors and the direct parent of a Node, respectively. - Made Template.remove(keep_field=True) behave more reasonably when the parameter is already empty. - Added the keep_template_params argument to Wikicode.strip_code(). If True, diff --git a/docs/changelog.rst b/docs/changelog.rst index 2c6be16..4d0d6fd 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -7,6 +7,11 @@ v0.5 Unreleased (`changes `__): +- Added :meth:`.Wikicode.contains` to determine whether a :class:`.Node` or + :class:`.Wikicode` object is contained within another :class:`.Wikicode` + object. +- Added :meth:`.Wikicode.get_ancestors` and :meth:`.Wikicode.get_parent` to + find all ancestors and the direct parent of a :class:`.Node`, respectively. - Made :meth:`Template.remove(keep_field=True) <.Template.remove>` behave more reasonably when the parameter is already empty. - Added the *keep_template_params* argument to :meth:`.Wikicode.strip_code`. diff --git a/mwparserfromhell/wikicode.py b/mwparserfromhell/wikicode.py index 73aea41..4379b0a 100644 --- a/mwparserfromhell/wikicode.py +++ b/mwparserfromhell/wikicode.py @@ -275,6 +275,21 @@ class Wikicode(StringMixIn): else: self.nodes.pop(index) + def contains(self, obj): + """Return whether this Wikicode object contains *obj*. + + If *obj* is a :class:`.Node` or :class:`.Wikicode` object, then we + search for it exactly among all of our children, recursively. + Otherwise, this method just uses :meth:`.__contains__` on the string. + """ + if not isinstance(obj, (Node, Wikicode)): + return obj in self + try: + self._do_strong_search(obj, recursive=True) + except ValueError: + return False + return True + def index(self, obj, recursive=False): """Return the index of *obj* in the list of nodes. @@ -294,6 +309,52 @@ class Wikicode(StringMixIn): return i raise ValueError(obj) + def get_ancestors(self, obj): + """Return a list of all ancestor nodes of the :class:`.Node` *obj*. + + The list is ordered from the most shallow ancestor (greatest great- + grandparent) to the direct parent. The node itself is not included in + the list. For example:: + + >>> text = "{{a|{{b|{{c|{{d}}}}}}}}" + >>> code = mwparserfromhell.parse(text) + >>> node = code.filter_templates(matches=lambda n: n == "{{d}}")[0] + >>> code.get_ancestors(node) + ['{{a|{{b|{{c|{{d}}}}}}}}', '{{b|{{c|{{d}}}}}}', '{{c|{{d}}}}'] + + Will return an empty list if *obj* is at the top level of this Wikicode + object. Will raise :exc:`ValueError` if it wasn't found. + """ + def _get_ancestors(code, needle): + for node in code.nodes: + if node is needle: + return [] + for code in node.__children__(): + ancestors = _get_ancestors(code, needle) + if ancestors is not None: + return [node] + ancestors + + if isinstance(obj, Wikicode): + obj = obj.get(0) + elif not isinstance(obj, Node): + raise ValueError(obj) + + ancestors = _get_ancestors(self, obj) + if ancestors is None: + raise ValueError(obj) + return ancestors + + def get_parent(self, obj): + """Return the direct parent node of the :class:`.Node` *obj*. + + This function is equivalent to calling :meth:`.get_ancestors` and + taking the last element of the resulting list. Will return None if + the node exists but does not have a parent; i.e., it is at the top + level of the Wikicode object. + """ + ancestors = self.get_ancestors(obj) + return ancestors[-1] if ancestors else None + def insert(self, index, value): """Insert *value* at *index* in the list of nodes. diff --git a/tests/test_wikicode.py b/tests/test_wikicode.py index 5457920..c77fdd2 100644 --- a/tests/test_wikicode.py +++ b/tests/test_wikicode.py @@ -85,6 +85,17 @@ class TestWikicode(TreeEqualityTestCase): self.assertRaises(IndexError, code.set, 3, "{{baz}}") self.assertRaises(IndexError, code.set, -4, "{{baz}}") + def test_contains(self): + """test Wikicode.contains()""" + code = parse("Here is {{aaa|{{bbb|xyz{{ccc}}}}}} and a [[page|link]]") + tmpl1, tmpl2, tmpl3 = code.filter_templates() + tmpl4 = parse("{{ccc}}").filter_templates()[0] + self.assertTrue(code.contains(tmpl1)) + self.assertTrue(code.contains(tmpl3)) + self.assertFalse(code.contains(tmpl4)) + self.assertTrue(code.contains(str(tmpl4))) + self.assertTrue(code.contains(tmpl2.params[0].value)) + def test_index(self): """test Wikicode.index()""" code = parse("Have a {{template}} and a [[page|link]]") @@ -102,6 +113,22 @@ class TestWikicode(TreeEqualityTestCase): self.assertRaises(ValueError, code.index, code.get(1).get(1).value, recursive=False) + def test_get_ancestors_parent(self): + """test Wikicode.get_ancestors() and Wikicode.get_parent()""" + code = parse("{{a|{{b|{{d|{{e}}{{f}}}}{{g}}}}}}{{c}}") + tmpl = code.filter_templates(matches=lambda n: n.name == "f")[0] + parent1 = code.filter_templates(matches=lambda n: n.name == "d")[0] + parent2 = code.filter_templates(matches=lambda n: n.name == "b")[0] + parent3 = code.filter_templates(matches=lambda n: n.name == "a")[0] + fake = parse("{{f}}").get(0) + + self.assertEqual([parent3, parent2, parent1], code.get_ancestors(tmpl)) + self.assertIs(parent1, code.get_parent(tmpl)) + self.assertEqual([], code.get_ancestors(parent3)) + self.assertIs(None, code.get_parent(parent3)) + self.assertRaises(ValueError, code.get_ancestors, fake) + self.assertRaises(ValueError, code.get_parent, fake) + def test_insert(self): """test Wikicode.insert()""" code = parse("Have a {{template}} and a [[page|link]]")