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

1 files changed, 47 insertions, 9 deletions

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)