Cleanup Wikicode's filter functions; implement test_filter_family().

mwparserfromhell/wikicode.py

@@ -24,7 +24,8 @@ from __future__ import unicode_literals
 import re
 
 from .compat import maxsize, py3k, str
-from .nodes import Heading, Node, Tag, Template, Text, Wikilink
+from .nodes import (Argument, Comment, Heading, HTMLEntity, Node, Tag,
+                    Template, Text, Wikilink)
 from .string_mixin import StringMixIn
 from .utils import parse_anything
 
@@ -151,6 +152,36 @@ class Wikicode(StringMixIn):
             node.__showtree__(write, get, mark)
         return lines
 
+    @classmethod
+    def _build_filter_methods(cls, **meths):
+        """Given Node types, build the corresponding i?filter shortcuts.
+
+        The should be given as keys storing the method's base name paired
+        with values storing the corresponding :py:class:`~.Node` type. For
+        example, the dict may contain the pair ``("templates", Template)``,
+        which will produce the methods :py:meth:`ifilter_templates` and
+        :py:meth:`filter_templates`, which are shortcuts for
+        :py:meth:`ifilter(forcetype=Template) <ifilter>` and
+        :py:meth:`filter(forcetype=Template) <filter>`, respectively. These
+        shortcuts are added to the class itself, with an appropriate docstring.
+        """
+        doc = """Iterate over {0}.
+
+        This is equivalent to :py:meth:`{1}` with *forcetype* set to
+        :py:class:`~.{2}`.
+        """
+        make_ifilter = lambda ftype: (lambda self, **kw:
+                                      self.ifilter(forcetype=ftype, **kw))
+        make_filter = lambda ftype: (lambda self, **kw:
+                                     self.filter(forcetype=ftype, **kw))
+        for name, ftype in (meths.items() if py3k else meths.iteritems()):
+            ifilter = make_ifilter(ftype)
+            filter = make_filter(ftype)
+            ifilter.__doc__ = doc.format(name, "ifilter", ftype.__name__)
+            filter.__doc__ = doc.format(name, "filter", ftype.__name__)
+            setattr(cls, "ifilter_" + name, ifilter)
+            setattr(cls, "filter_" + name, filter)
+
     @property
     def nodes(self):
         """A list of :py:class:`~.Node` objects.
@@ -296,32 +327,6 @@ class Wikicode(StringMixIn):
                 if not matches or re.search(matches, str(node), flags):
                     yield node
 
-    @classmethod
-    def _build_filter_methods(cls, meths):
-        """Given a dict of Node types, build corresponding i?filter shortcuts.
-
-        The dict should be given as keys storing the method's base name paired
-        with values storing the corresponding :py:class:`~.Node` type. For
-        example, the dict may contain the pair ``("templates", Template)``,
-        which will produce the methods :py:meth:`ifilter_templates` and
-        :py:meth:`filter_templates`, which are shortcuts for
-        :py:meth:`ifilter(forcetype=Template) <ifilter>` and
-        :py:meth:`filter(forcetype=Template) <filter>`, respectively. These
-        shortcuts are added to the class itself, with an appropriate docstring.
-        """
-        doc = """Iterate over {0}.
-
-        This is equivalent to :py:meth:`{1}` with *forcetype* set to
-        :py:class:`~.{2}`.
-        """
-        for name, forcetype in (meths.items() if py3k else meths.iteritems()):
-            ifil = lambda self, **kw: self.ifilter(forcetype=forcetype, **kw)
-            fil = lambda self, **kw: self.filter(forcetype=forcetype, **kw)
-            ifil.__doc__ = doc.format(name, "ifilter", forcetype)
-            fil.__doc__ = doc.format(name, "filter", forcetype)
-            setattr(cls, "ifilter_" + name, ifil)
-            setattr(cls, "filter_" + name, fil)
-
     def filter(self, recursive=False, matches=None, flags=FLAGS,
                forcetype=None):
         """Return a list of nodes within our list matching certain conditions.
@@ -429,9 +434,7 @@ class Wikicode(StringMixIn):
         marker = object()  # Random object we can find with certainty in a list
         return "\n".join(self._get_tree(self, [], marker, 0))
 
-Wikicode._build_filter_methods({
-    "links": Wikilink,
-    "templates": Template,
-    "text": Text,
-    "tag": Tag
-    })
+Wikicode._build_filter_methods(
+    arguments=Argument, comments=Comment, headings=Heading,
+    html_entities=HTMLEntity, tags=Tag, templates=Template, text=Text,
+    wikilinks=Wikilink)

tests/test_wikicode.py

@@ -21,6 +21,8 @@
 # SOFTWARE.
 
 from __future__ import unicode_literals
+import re
+from types import GeneratorType
 import unittest
 
 from mwparserfromhell.nodes import (Argument, Comment, Heading, HTMLEntity,
@@ -210,7 +212,67 @@ class TestWikicode(TreeEqualityTestCase):
 
     def test_filter_family(self):
         """test the Wikicode.i?filter() family of functions"""
-        pass
+        def genlist(gen):
+            self.assertIsInstance(gen, GeneratorType)
+            return list(gen)
+        ifilter = lambda code: (lambda **kw: genlist(code.ifilter(**kw)))
+
+        code = parse("a{{b}}c[[d]]{{{e}}}{{f}}[[g]]")
+        for func in (code.filter, ifilter(code)):
+            self.assertEqual(["a", "{{b}}", "c", "[[d]]", "{{{e}}}", "{{f}}",
+                              "[[g]]"], func())
+            self.assertEqual(["{{{e}}}"], func(forcetype=Argument))
+            self.assertIs(code.get(4), func(forcetype=Argument)[0])
+            self.assertEqual(["a", "c"], func(forcetype=Text))
+            self.assertEqual([], func(forcetype=Heading))
+            self.assertRaises(TypeError, func, forcetype=True)
+
+        funcs = [
+            lambda name, **kw: getattr(code, "filter_" + name)(**kw),
+            lambda name, **kw: genlist(getattr(code, "ifilter_" + name)(**kw))
+        ]
+        for get_filter in funcs:
+            self.assertEqual(["{{{e}}}"], get_filter("arguments"))
+            self.assertIs(code.get(4), get_filter("arguments")[0])
+            self.assertEqual([], get_filter("comments"))
+            self.assertEqual([], get_filter("headings"))
+            self.assertEqual([], get_filter("html_entities"))
+            self.assertEqual([], get_filter("tags"))
+            self.assertEqual(["{{b}}", "{{f}}"], get_filter("templates"))
+            self.assertEqual(["a", "c"], get_filter("text"))
+            self.assertEqual(["[[d]]", "[[g]]"], get_filter("wikilinks"))
+
+        code2 = parse("{{a|{{b}}|{{c|d={{f}}{{h}}}}}}")
+        for func in (code2.filter, ifilter(code2)):
+            self.assertEqual(["{{a|{{b}}|{{c|d={{f}}{{h}}}}}}"],
+                             func(recursive=False, forcetype=Template))
+            self.assertEqual(["{{a|{{b}}|{{c|d={{f}}{{h}}}}}}", "{{b}}",
+                              "{{c|d={{f}}{{h}}}}", "{{f}}", "{{h}}"],
+                             func(recursive=True, forcetype=Template))
+
+        code3 = parse("{{foobar}}{{FOO}}{{baz}}{{bz}}")
+        for func in (code3.filter, ifilter(code3)):
+            self.assertEqual(["{{foobar}}", "{{FOO}}"], func(matches=r"foo"))
+            self.assertEqual(["{{foobar}}", "{{FOO}}"],
+                             func(matches=r"^{{foo.*?}}"))
+            self.assertEqual(["{{foobar}}"],
+                             func(matches=r"^{{foo.*?}}", flags=re.UNICODE))
+            self.assertEqual(["{{baz}}", "{{bz}}"], func(matches=r"^{{b.*?z"))
+            self.assertEqual(["{{baz}}"], func(matches=r"^{{b.+?z}}"))
+
+        self.assertEqual(["{{a|{{b}}|{{c|d={{f}}{{h}}}}}}"],
+                         code2.filter_templates(recursive=False))
+        self.assertEqual(["{{a|{{b}}|{{c|d={{f}}{{h}}}}}}", "{{b}}",
+                          "{{c|d={{f}}{{h}}}}", "{{f}}", "{{h}}"],
+                         code2.filter_templates(recursive=True))
+        self.assertEqual(["{{baz}}", "{{bz}}"],
+                         code3.filter_templates(matches=r"^{{b.*?z"))
+        self.assertEqual([], code3.filter_tags(matches=r"^{{b.*?z"))
+        self.assertEqual([], code3.filter_tags(matches=r"^{{b.*?z", flags=0))
+
+        self.assertRaises(TypeError, code.filter_templates, 100)
+        self.assertRaises(TypeError, code.filter_templates, a=42)
+        self.assertRaises(TypeError, code.filter_templates, forcetype=Template)
 
     def test_get_sections(self):
         """test Wikicode.get_sections()"""
@@ -224,6 +286,5 @@ class TestWikicode(TreeEqualityTestCase):
         """test Wikicode.get_tree()"""
         pass
 
-
 if __name__ == "__main__":
     unittest.main(verbosity=2)