Simplify fix_zip

Part of #72

Author: brettcannon

libmodernize/fixes/fix_zip.py

@@ -1,40 +1,17 @@
 # Copyright 2008 Armin Ronacher.
 # Licensed to PSF under a Contributor Agreement.
 
-from lib2to3 import fixer_base
-from lib2to3.fixer_util import Name, Call, touch_import, in_special_context
+from lib2to3 import fixer_util
+from lib2to3.fixes import fix_zip
 
 
-class FixZip(fixer_base.ConditionalFix):
+class FixZip(fix_zip.FixZip):
 
-    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.
+        fixer_util.touch_import(u'six.moves', u'zip', node)
+        return result

tests/test_fix_zip.py

@@ -31,12 +31,6 @@
 list(zip(*args))
 """)
 
-ZIP_KWARGS = ("""\
-zip(arg1=[1])
-""", """\
-zip(arg1=[1])
-""")
-
 ZIP_ITERATOR_CONTEXT = ("""\
 for a in zip(x):
     pass
@@ -59,8 +53,5 @@ def test_zip_call_2_args():
 def test_zip_call_star_args():
     check_on_input(*ZIP_CALL_STAR_ARGS)
 
-def test_zip_kwargs():
-    check_on_input(*ZIP_KWARGS)
-
 def test_zip_iterator_context():
     check_on_input(*ZIP_ITERATOR_CONTEXT)