bpo-28837: Fix lib2to3 handling of map/zip/filter calls when followed with a 'trailer', e.g. zip()[x] (#24)

4 files changed, 109 insertions, 26 deletions

diff --git a/Lib/lib2to3/fixes/fix_filter.py b/Lib/lib2to3/fixes/fix_filter.py
index bb6718c..a7a5a15 100644
--- a/Lib/lib2to3/fixes/fix_filter.py
+++ b/Lib/lib2to3/fixes/fix_filter.py
@@ -15,7 +15,10 @@ Python 2.6 figure it out.
 
 # Local imports
 from .. import fixer_base
-from ..fixer_util import Name, Call, ListComp, in_special_context
+from ..pytree import Node
+from ..pygram import python_symbols as syms
+from ..fixer_util import Name, ArgList, ListComp, in_special_context
+
 
 class FixFilter(fixer_base.ConditionalFix):
     BM_compatible = True
@@ -34,16 +37,19 @@ class FixFilter(fixer_base.ConditionalFix):
             >
             ')'
         >
+        [extra_trailers=trailer*]
     >
     |
     power<
         'filter'
         trailer< '(' arglist< none='None' ',' seq=any > ')' >
+        [extra_trailers=trailer*]
     >
     |
     power<
         'filter'
         args=trailer< '(' [any] ')' >
+        [extra_trailers=trailer*]
     >
     """
 
@@ -53,23 +59,32 @@ class FixFilter(fixer_base.ConditionalFix):
         if self.should_skip(node):
             return
 
+        trailers = []
+        if 'extra_trailers' in results:
+            for t in results['extra_trailers']:
+                trailers.append(t.clone())
+
         if "filter_lambda" in results:
             new = ListComp(results.get("fp").clone(),
                            results.get("fp").clone(),
                            results.get("it").clone(),
                            results.get("xp").clone())
+            new = Node(syms.power, [new] + trailers, prefix="")
 
         elif "none" in results:
             new = ListComp(Name("_f"),
                            Name("_f"),
                            results["seq"].clone(),
                            Name("_f"))
+            new = Node(syms.power, [new] + trailers, prefix="")
 
         else:
             if in_special_context(node):
                 return None
-            new = node.clone()
+
+            args = results['args'].clone()
+            new = Node(syms.power, [Name("filter"), args], prefix="")
+            new = Node(syms.power, [Name("list"), ArgList([new])] + trailers)
             new.prefix = ""
-            new = Call(Name("list"), [new])
         new.prefix = node.prefix
         return new
diff --git a/Lib/lib2to3/fixes/fix_map.py b/Lib/lib2to3/fixes/fix_map.py
index 9f966fe..78cf81c 100644
--- a/Lib/lib2to3/fixes/fix_map.py
+++ b/Lib/lib2to3/fixes/fix_map.py
@@ -22,8 +22,10 @@ soon as the shortest argument is exhausted.
 # Local imports
 from ..pgen2 import token
 from .. import fixer_base
-from ..fixer_util import Name, Call, ListComp, in_special_context
+from ..fixer_util import Name, ArgList, Call, ListComp, in_special_context
 from ..pygram import python_symbols as syms
+from ..pytree import Node
+
 
 class FixMap(fixer_base.ConditionalFix):
     BM_compatible = True
@@ -32,6 +34,7 @@ class FixMap(fixer_base.ConditionalFix):
     map_none=power<
         'map'
         trailer< '(' arglist< 'None' ',' arg=any [','] > ')' >
+        [extra_trailers=trailer*]
     >
     |
     map_lambda=power<
@@ -47,10 +50,12 @@ class FixMap(fixer_base.ConditionalFix):
             >
             ')'
         >
+        [extra_trailers=trailer*]
     >
     |
     power<
-        'map' trailer< '(' [arglist=any] ')' >
+        'map' args=trailer< '(' [any] ')' >
+        [extra_trailers=trailer*]
     >
     """
 
@@ -60,6 +65,11 @@ class FixMap(fixer_base.ConditionalFix):
         if self.should_skip(node):
             return
 
+        trailers = []
+        if 'extra_trailers' in results:
+            for t in results['extra_trailers']:
+                trailers.append(t.clone())
+
         if node.parent.type == syms.simple_stmt:
             self.warning(node, "You should use a for loop here")
             new = node.clone()
@@ -69,23 +79,32 @@ class FixMap(fixer_base.ConditionalFix):
             new = ListComp(results["xp"].clone(),
                            results["fp"].clone(),
                            results["it"].clone())
+            new = Node(syms.power, [new] + trailers, prefix="")
+
         else:
             if "map_none" in results:
                 new = results["arg"].clone()
+                new.prefix = ""
             else:
-                if "arglist" in results:
-                    args = results["arglist"]
-                    if args.type == syms.arglist and \
-                       args.children[0].type == token.NAME and \
-                       args.children[0].value == "None":
+                if "args" in results:
+                    args = results["args"]
+                    if args.type == syms.trailer and \
+                       args.children[1].type == syms.arglist and \
+                       args.children[1].children[0].type == token.NAME and \
+                       args.children[1].children[0].value == "None":
                         self.warning(node, "cannot convert map(None, ...) "
                                      "with multiple arguments because map() "
                                      "now truncates to the shortest sequence")
                         return
+
+                    new = Node(syms.power, [Name("map"), args.clone()])
+                    new.prefix = ""
+
                 if in_special_context(node):
                     return None
-                new = node.clone()
+
+            new = Node(syms.power, [Name("list"), ArgList([new])] + trailers)
             new.prefix = ""
-            new = Call(Name("list"), [new])
+
         new.prefix = node.prefix
         return new
diff --git a/Lib/lib2to3/fixes/fix_zip.py b/Lib/lib2to3/fixes/fix_zip.py
index 8f36a94..52c28df 100644
--- a/Lib/lib2to3/fixes/fix_zip.py
+++ b/Lib/lib2to3/fixes/fix_zip.py
@@ -9,13 +9,16 @@ iter(<>), list(<>), tuple(<>), sorted(<>), ...join(<>), or for V in <>:.
 
 # Local imports
 from .. import fixer_base
-from ..fixer_util import Name, Call, in_special_context
+from ..pytree import Node
+from ..pygram import python_symbols as syms
+from ..fixer_util import Name, ArgList, in_special_context
+
 
 class FixZip(fixer_base.ConditionalFix):
 
     BM_compatible = True
     PATTERN = """
-    power< 'zip' args=trailer< '(' [any] ')' >
+    power< 'zip' args=trailer< '(' [any] ')' > [trailers=trailer*]
     >
     """
 
@@ -28,8 +31,16 @@ class FixZip(fixer_base.ConditionalFix):
         if in_special_context(node):
             return None
 
-        new = node.clone()
-        new.prefix = ""
-        new = Call(Name("list"), [new])
+        args = results['args'].clone()
+        args.prefix = ""
+
+        trailers = []
+        if 'trailers' in results:
+            trailers = [n.clone() for n in results['trailers']]
+            for n in trailers:
+                n.prefix = ""
+
+        new = Node(syms.power, [Name("zip"), args], prefix="")
+        new = Node(syms.power, [Name("list"), ArgList([new])] + trailers)
         new.prefix = node.prefix
         return new
diff --git a/Lib/lib2to3/tests/test_fixers.py b/Lib/lib2to3/tests/test_fixers.py
index b3f2680..3e1a255 100644
--- a/Lib/lib2to3/tests/test_fixers.py
+++ b/Lib/lib2to3/tests/test_fixers.py
@@ -2954,10 +2954,23 @@ class Test_filter(FixerTestCase):
         a = """x = [x for x in range(10) if x%2 == 0]"""
         self.check(b, a)
 
-        # XXX This (rare) case is not supported
-##         b = """x = filter(f, 'abc')[0]"""
-##         a = """x = list(filter(f, 'abc'))[0]"""
-##         self.check(b, a)
+    def test_filter_trailers(self):
+        b = """x = filter(None, 'abc')[0]"""
+        a = """x = [_f for _f in 'abc' if _f][0]"""
+        self.check(b, a)
+
+        b = """x = len(filter(f, 'abc')[0])"""
+        a = """x = len(list(filter(f, 'abc'))[0])"""
+        self.check(b, a)
+
+        b = """x = filter(lambda x: x%2 == 0, range(10))[0]"""
+        a = """x = [x for x in range(10) if x%2 == 0][0]"""
+        self.check(b, a)
+
+        # Note the parens around x
+        b = """x = filter(lambda (x): x%2 == 0, range(10))[0]"""
+        a = """x = [x for x in range(10) if x%2 == 0][0]"""
+        self.check(b, a)
 
     def test_filter_nochange(self):
         a = """b.join(filter(f, 'abc'))"""
@@ -3022,6 +3035,23 @@ class Test_map(FixerTestCase):
         a = """x =    list(map(   f,    'abc'   ))"""
         self.check(b, a)
 
+    def test_map_trailers(self):
+        b = """x = map(f, 'abc')[0]"""
+        a = """x = list(map(f, 'abc'))[0]"""
+        self.check(b, a)
+
+        b = """x = map(None, l)[0]"""
+        a = """x = list(l)[0]"""
+        self.check(b, a)
+
+        b = """x = map(lambda x:x, l)[0]"""
+        a = """x = [x for x in l][0]"""
+        self.check(b, a)
+
+        b = """x = map(f, 'abc')[0][1]"""
+        a = """x = list(map(f, 'abc'))[0][1]"""
+        self.check(b, a)
+
     def test_trailing_comment(self):
         b = """x = map(f, 'abc')   #   foo"""
         a = """x = list(map(f, 'abc'))   #   foo"""
@@ -3066,11 +3096,6 @@ class Test_map(FixerTestCase):
             """
         self.warns(b, a, "You should use a for loop here")
 
-        # XXX This (rare) case is not supported
-##         b = """x = map(f, 'abc')[0]"""
-##         a = """x = list(map(f, 'abc'))[0]"""
-##         self.check(b, a)
-
     def test_map_nochange(self):
         a = """b.join(map(f, 'abc'))"""
         self.unchanged(a)
@@ -3130,6 +3155,10 @@ class Test_zip(FixerTestCase):
         super(Test_zip, self).check(b, a)
 
     def test_zip_basic(self):
+        b = """x = zip()"""
+        a = """x = list(zip())"""
+        self.check(b, a)
+
         b = """x = zip(a, b, c)"""
         a = """x = list(zip(a, b, c))"""
         self.check(b, a)
@@ -3138,6 +3167,15 @@ class Test_zip(FixerTestCase):
         a = """x = len(list(zip(a, b)))"""
         self.check(b, a)
 
+    def test_zip_trailers(self):
+        b = """x = zip(a, b, c)[0]"""
+        a = """x = list(zip(a, b, c))[0]"""
+        self.check(b, a)
+
+        b = """x = zip(a, b, c)[0][1]"""
+        a = """x = list(zip(a, b, c))[0][1]"""
+        self.check(b, a)
+
     def test_zip_nochange(self):
         a = """b.join(zip(a, b))"""
         self.unchanged(a)