Simplify fix_filter

Part of #72

Author: brettcannon

libmodernize/fixes/fix_filter.py

@@ -1,42 +1,17 @@
 # Copyright 2008 Armin Ronacher.
 # Licensed to PSF under a Contributor Agreement.
 
-from lib2to3 import fixer_base
-from lib2to3.fixer_util import Call, Name, touch_import, in_special_context
+from lib2to3 import fixer_util
+from lib2to3.fixes import fix_filter
 
 
-class FixFilter(fixer_base.ConditionalFix):
+class FixFilter(fix_filter.FixFilter):
 
-    BM_compatible = True
-    order = "pre"
     skip_on = "six.moves.filter"
 
-    PATTERN = """
-    power< 'filter'
-        trailer< '('
-            arglist< (
-                not(argument<any '=' any>) any ','
-                not(argument<any '=' any>) any
-            ) >
-        ')' >
-    >
-    """
-
     def transform(self, node, results):
-        if self.should_skip(node):
-            # Make the fixer idempotent - if six.moves.filter is already imported,
-            # skip it. should_skip() caches the state the first time we check,
-            # so it doesn't skip after *we've* added the six.moves import.
-            return
-
-        touch_import(u'six.moves', u'filter', node)
-        if in_special_context(node):
-            # The node is somewhere where it only needs to be an iterator,
-            # e.g. a for loop - don't wrap in list()
-            return
-
-        new = node.clone()
-        new.prefix = ""
-        new = Call(Name("list"), [new])
-        new.prefix = node.prefix
-        return new
+        result = super(FixFilter, self).transform(node, results)
+        # Keep performance improvement from six.moves.filter in iterator
+        # contexts on Python 2.7.
+        fixer_util.touch_import(u'six.moves', u'filter', node)
+        return result

tests/test_fix_filter.py

@@ -2,51 +2,25 @@
 
 
 FILTER_CALL = ("""\
-filter(None, [1])
+filter(func, [1])
 """, """\
 from six.moves import filter
-list(filter(None, [1]))
+list(filter(func, [1]))
 """)
 
-FILTER_TOO_FEW_ARGS = ("""\
-filter(None)
-""", """\
-filter(None)
-""")
-
-FILTER_TOO_MANY_ARGS = ("""\
-filter(None, [1], [2])
-""", """\
-filter(None, [1], [2])
-""")
-
-FILTER_KWARGS = ("""\
-filter(function=None, [1])
-""", """\
-filter(function=None, [1])
-""")
 
 FILTER_ITERATOR_CONTEXT = ("""\
-for a in filter(None, [1]):
+for a in filter(func, [1]):
     pass
 """, """\
 from six.moves import filter
-for a in filter(None, [1]):
+for a in filter(func, [1]):
     pass
 """)
 
 
 def test_filter_call():
     check_on_input(*FILTER_CALL)
 
-def test_filter_too_few_args():
-    check_on_input(*FILTER_TOO_FEW_ARGS)
-
-def test_filter_too_many_args():
-    check_on_input(*FILTER_TOO_MANY_ARGS)
-
-def test_filter_kwargs():
-    check_on_input(*FILTER_KWARGS)
-
 def test_filter_iterator_context():
     check_on_input(*FILTER_ITERATOR_CONTEXT)