Add Wikicode.contains(), Wikicode.get_ancestors(), Wikicode.get_parent() (#177)

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,

docs/changelog.rst

@@ -7,6 +7,11 @@ v0.5
 Unreleased
 (`changes <https://github.com/earwig/mwparserfromhell/compare/v0.4.4...develop>`__):
 
+- 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`.

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.
 

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]]")