Merge pull request #75 from brettcannon/simplification

Simplify various fixers

Author: takluyver

libmodernize/__init__.py

@@ -67,3 +67,13 @@ def add_future(node, symbol):
 def touch_import(package, name, node):
     add_future(node, 'absolute_import')
     fixer_util.touch_import(package, name, node)
+
+
+def is_listcomp(node):
+    return (isinstance(node, Node) and
+             node.type == syms.atom and
+             len(node.children) >= 2 and
+             isinstance(node.children[0], Leaf) and
+             node.children[0].value == '[' and
+             isinstance(node.children[-1], Leaf) and
+             node.children[-1].value == ']')

libmodernize/fixes/fix_filter.py

@@ -2,43 +2,18 @@
 # Licensed to PSF under a Contributor Agreement.
 from __future__ import absolute_import
 
-from lib2to3 import fixer_base
-from lib2to3.fixer_util import Call, Name, in_special_context
-from libmodernize import touch_import
+from lib2to3.fixes import fix_filter
+import libmodernize
 
 
-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)
+        if not libmodernize.is_listcomp(result):
+            # Keep performance improvement from six.moves.filter in iterator
+            # contexts on Python 2.7.
+            libmodernize.touch_import(u'six.moves', u'filter', node)
+        return result

libmodernize/fixes/fix_map.py

@@ -2,43 +2,17 @@
 # Licensed to PSF under a Contributor Agreement.
 from __future__ import absolute_import
 
-from lib2to3 import fixer_base
-from lib2to3.fixer_util import Call, Name, in_special_context
-from libmodernize import touch_import
+from lib2to3.fixes import fix_map
+import libmodernize
 
+class FixMap(fix_map.FixMap):
 
-class FixMap(fixer_base.ConditionalFix):
-
-    BM_compatible = True
-    order = "pre"
     skip_on = "six.moves.map"
 
-    PATTERN = """
-    power< 'map'
-        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.map 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'map', 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(FixMap, self).transform(node, results)
+        if not libmodernize.is_listcomp(result):
+            # Always use the import even if no change is required so as to have
+            # improved performance in iterator contexts even on Python 2.7.
+            libmodernize.touch_import(u'six.moves', u'map', node)
+        return result

libmodernize/fixes/fix_print.py

@@ -1,92 +1,13 @@
 # Copyright 2006 Google, Inc. All Rights Reserved.
 # Licensed to PSF under a Contributor Agreement.
 
-"""Fixer for print.
+from lib2to3.fixes import fix_print
+import libmodernize
 
-Change:
-    "print"          into "print()"
-    "print ..."      into "print(...)"
-    "print(...)"     not changed
-    "print ... ,"    into "print(..., end=' ')"
-    "print >>x, ..." into "print(..., file=x)"
 
-No changes are applied if print_function is imported from __future__
-
-"""
-from __future__ import absolute_import
-
-# Local imports
-from lib2to3 import patcomp, pytree, fixer_base
-from lib2to3.pgen2 import token
-from lib2to3.fixer_util import Name, Call, Comma, String
-from libmodernize import add_future
-
-parend_expr = patcomp.compile_pattern(
-              """atom< '(' [arith_expr|atom|power|term|STRING|NAME] ')' >"""
-              )
-
-
-class FixPrint(fixer_base.BaseFix):
-
-    BM_compatible = True
-
-    PATTERN = """
-              simple_stmt< any* bare='print' any* > | print_stmt
-              """
+class FixPrint(fix_print.FixPrint):
 
     def transform(self, node, results):
-        assert results
-
-        bare_print = results.get("bare")
-
-        if bare_print:
-            # Special-case print all by itself.
-            bare_print.replace(Call(Name(u"print"), [],
-                               prefix=bare_print.prefix))
-            add_future(node, u'print_function')
-            return
-        assert node.children[0] == Name(u"print")
-        args = node.children[1:]
-        if len(args) == 1 and parend_expr.match(args[0]):
-            # We don't want to keep sticking parens around an
-            # already-parenthesised expression.
-            return
-
-        end = file = None
-        if args and args[-1] == Comma():
-            args = args[:-1]
-            end = " "
-        if args and args[0] == pytree.Leaf(token.RIGHTSHIFT, u">>"):
-            assert len(args) >= 2
-            file = args[1].clone()
-            args = args[3:] # Strip a possible comma after the file expression
-        # Now synthesize a print(args, end=..., file=...) node
-        # (sep is not used).
-        l_args = [arg.clone() for arg in args]
-        if l_args:
-            l_args[0].prefix = u""
-        if end is not None or file is not None:
-            if end is not None:
-                self.add_kwarg(l_args, u"end", String(repr(end)))
-            if file is not None:
-                self.add_kwarg(l_args, u"file", file)
-        n_stmt = Call(Name(u"print"), l_args)
-        n_stmt.prefix = node.prefix
-
-        # Note that there are corner cases where adding this future-import is
-        # incorrect, for example when the file also has a 'print ()' statement
-        # that was intended to print "()".
-        add_future(node, u'print_function')
-        return n_stmt
-
-    def add_kwarg(self, l_nodes, s_kwd, n_expr):
-        # XXX All this prefix-setting may lose comments (though rarely)
-        n_expr.prefix = u""
-        n_argument = pytree.Node(self.syms.argument,
-                                 (Name(s_kwd),
-                                  pytree.Leaf(token.EQUAL, u"="),
-                                  n_expr))
-        if l_nodes:
-            l_nodes.append(Comma())
-            n_argument.prefix = u" "
-        l_nodes.append(n_argument)
+        result = super(FixPrint, self).transform(node, results)
+        libmodernize.add_future(node, u'print_function')
+        return result

libmodernize/fixes/fix_zip.py

@@ -2,41 +2,16 @@
 # Licensed to PSF under a Contributor Agreement.
 from __future__ import absolute_import
 
-from lib2to3 import fixer_base
-from lib2to3.fixer_util import Name, Call, in_special_context
-from libmodernize import touch_import
+from lib2to3.fixes import fix_zip
+import libmodernize
 
+class FixZip(fix_zip.FixZip):
 
-class FixZip(fixer_base.ConditionalFix):
-
-    BM_compatible = True
-    order = "pre"
     skip_on = "six.moves.zip"
 
-    PATTERN = """
-    power< 'zip'
-        trailer< '('
-            ( not(arglist | argument<any '=' any>) any* |
-              arglist < not(argument<any '=' any>) any* > )
-        ')' >
-    >
-    """
-
     def transform(self, node, results):
-        if self.should_skip(node):
-            # Make the fixer idempotent - if six.moves.zip 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'zip', 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(FixZip, self).transform(node, results)
+        # Always use six.moves.zip so that even Python 2.7 gets performance
+        # boost from using itertools in iterator contexts.
+        libmodernize.touch_import(u'six.moves', u'zip', node)
+        return result

tests/test_fix_filter.py

@@ -4,53 +4,36 @@
 
 
 FILTER_CALL = ("""\
-filter(None, [1])
+filter(func, [1])
 """, """\
 from __future__ import absolute_import
 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 __future__ import absolute_import
 from six.moves import filter
-for a in filter(None, [1]):
+for a in filter(func, [1]):
     pass
 """)
 
+FILTER_NONE = ("""\
+filter(None, x)
+""", """\
+[_f for _f in x if _f]
+""")
+
 
 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)
+
+def test_filter_None():
+    check_on_input(*FILTER_NONE)