Skip six imports for map and filter when changed to a list comp

Author: brettcannon

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,6 +2,9 @@
 # Licensed to PSF under a Contributor Agreement.
 from __future__ import absolute_import
 
+from lib2to3 import fixer_util
+from lib2to3 import pygram
+from lib2to3 import pytree
 from lib2to3.fixes import fix_filter
 import libmodernize
 
@@ -12,7 +15,8 @@ class FixFilter(fix_filter.FixFilter):
 
     def transform(self, node, results):
         result = super(FixFilter, self).transform(node, results)
-        # Keep performance improvement from six.moves.filter in iterator
-        # contexts on Python 2.7.
-        libmodernize.touch_import(u'six.moves', u'filter', node)
+        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,6 +2,7 @@
 # Licensed to PSF under a Contributor Agreement.
 from __future__ import absolute_import
 
+from lib2to3 import fixer_util
 from lib2to3.fixes import fix_map
 import libmodernize
 
@@ -11,7 +12,8 @@ class FixMap(fix_map.FixMap):
 
     def transform(self, node, results):
         result = super(FixMap, self).transform(node, results)
-        # 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)
+        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

tests/test_fix_filter.py

@@ -22,9 +22,18 @@
     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_iterator_context():
     check_on_input(*FILTER_ITERATOR_CONTEXT)
+
+def test_filter_None():
+    check_on_input(*FILTER_NONE)

tests/test_fix_map.py

@@ -50,6 +50,12 @@
     pass
 """)
 
+MAP_LAMBDA = ("""\
+x = map(lambda x: x+1, stuff)
+""", """\
+x = [x+1 for x in stuff]
+""")
+
 
 def test_map_1_arg():
     check_on_input(*MAP_1_ARG)
@@ -68,3 +74,6 @@ def test_map_ref():
 
 def test_map_iterator_context():
     check_on_input(*MAP_ITERATOR_CONTEXT)
+
+def test_map_lambda():
+    check_on_input(*MAP_LAMBDA)