Simplify fix_map

Part of #72

Author: brettcannon

libmodernize/fixes/fix_map.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_map
 
 
-class FixMap(fixer_base.ConditionalFix):
+class FixMap(fix_map.FixMap):
 
-    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)
+        # Always use the import even if no change is required so as to have
+        # improved performance in iterator contexts even on Python 2.7.
+        fixer_util.touch_import(u'six.moves', u'map', node)
+        return result

tests/test_fix_map.py

@@ -1,5 +1,11 @@
 from utils import check_on_input
 
+MAP_1_ARG = ("""\
+map(*args)
+""", """\
+from six.moves import map
+list(map(*args))
+""")
 
 MAP_2_ARGS = ("""\
 map(x, [1])
@@ -22,18 +28,6 @@
 list(map(x, [1], [2], [3]))
 """)
 
-MAP_TOO_FEW_ARGS = ("""\
-map(x)
-""", """\
-map(x)
-""")
-
-MAP_KWARGS = ("""\
-map(function=x, [1])
-""", """\
-map(function=x, [1])
-""")
-
 MAP_REF = ("""\
 x = map
 """, """\
@@ -50,6 +44,9 @@
 """)
 
 
+def test_map_1_arg():
+    check_on_input(*MAP_1_ARG)
+
 def test_map_2_args():
     check_on_input(*MAP_2_ARGS)
 
@@ -59,12 +56,6 @@ def test_map_3_args():
 def test_map_4_args():
     check_on_input(*MAP_4_ARGS)
 
-def test_map_too_few_args():
-    check_on_input(*MAP_TOO_FEW_ARGS)
-
-def test_map_kwargs():
-    check_on_input(*MAP_KWARGS)
-
 def test_map_ref():
     check_on_input(*MAP_REF)