[Python-checkins] (2.7): transform izip_longest #11424

[benjamin.peterson]

http://hg.python.org/cpython/rev/770d45d22a40
changeset:   68327:770d45d22a40
branch:      2.7
parent:      68321:d9cc58f93d72
user:        Benjamin Peterson <benjamin at python.org>
date:        Mon Mar 07 22:50:37 2011 -0600
summary:
  transform izip_longest #11424

files:
  Lib/lib2to3/fixes/fix_itertools.py
  Lib/lib2to3/fixes/fix_itertools_imports.py
  Lib/lib2to3/tests/test_fixers.py

diff --git a/Lib/lib2to3/fixes/fix_itertools.py b/Lib/lib2to3/fixes/fix_itertools.py
--- a/Lib/lib2to3/fixes/fix_itertools.py
+++ b/Lib/lib2to3/fixes/fix_itertools.py
@@ -13,7 +13,7 @@
 
 class FixItertools(fixer_base.BaseFix):
     BM_compatible = True
-    it_funcs = "('imap'|'ifilter'|'izip'|'ifilterfalse')"
+    it_funcs = "('imap'|'ifilter'|'izip'|'izip_longest'|'ifilterfalse')"
     PATTERN = """
               power< it='itertools'
                   trailer<
@@ -28,7 +28,8 @@
     def transform(self, node, results):
         prefix = None
         func = results['func'][0]
-        if 'it' in results and func.value != u'ifilterfalse':
+        if ('it' in results and
+            func.value not in (u'ifilterfalse', u'izip_longest')):
             dot, it = (results['dot'], results['it'])
             # Remove the 'itertools'
             prefix = it.prefix
diff --git a/Lib/lib2to3/fixes/fix_itertools_imports.py b/Lib/lib2to3/fixes/fix_itertools_imports.py
--- a/Lib/lib2to3/fixes/fix_itertools_imports.py
+++ b/Lib/lib2to3/fixes/fix_itertools_imports.py
@@ -31,9 +31,10 @@
             if member_name in (u'imap', u'izip', u'ifilter'):
                 child.value = None
                 child.remove()
-            elif member_name == u'ifilterfalse':
+            elif member_name in (u'ifilterfalse', u'izip_longest'):
                 node.changed()
-                name_node.value = u'filterfalse'
+                name_node.value = (u'filterfalse' if member_name[1] == u'f'
+                                   else u'zip_longest')
 
         # Make sure the import statement is still sane
         children = imports.children[:] or [imports]
diff --git a/Lib/lib2to3/tests/test_fixers.py b/Lib/lib2to3/tests/test_fixers.py
--- a/Lib/lib2to3/tests/test_fixers.py
+++ b/Lib/lib2to3/tests/test_fixers.py
@@ -3623,16 +3623,24 @@
         a = """%s(f, a)"""
         self.checkall(b, a)
 
-    def test_2(self):
+    def test_qualified(self):
         b = """itertools.ifilterfalse(a, b)"""
         a = """itertools.filterfalse(a, b)"""
         self.check(b, a)
 
-    def test_4(self):
+        b = """itertools.izip_longest(a, b)"""
+        a = """itertools.zip_longest(a, b)"""
+        self.check(b, a)
+
+    def test_2(self):
         b = """ifilterfalse(a, b)"""
         a = """filterfalse(a, b)"""
         self.check(b, a)
 
+        b = """izip_longest(a, b)"""
+        a = """zip_longest(a, b)"""
+        self.check(b, a)
+
     def test_space_1(self):
         b = """    %s(f, a)"""
         a = """    %s(f, a)"""
@@ -3643,9 +3651,14 @@
         a = """    itertools.filterfalse(a, b)"""
         self.check(b, a)
 
+        b = """    itertools.izip_longest(a, b)"""
+        a = """    itertools.zip_longest(a, b)"""
+        self.check(b, a)
+
     def test_run_order(self):
         self.assert_runs_after('map', 'zip', 'filter')
 
+
 class Test_itertools_imports(FixerTestCase):
     fixer = 'itertools_imports'
 
@@ -3696,18 +3709,19 @@
         s = "from itertools import bar as bang"
         self.unchanged(s)
 
-    def test_ifilter(self):
-        b = "from itertools import ifilterfalse"
-        a = "from itertools import filterfalse"
-        self.check(b, a)
-
-        b = "from itertools import imap, ifilterfalse, foo"
-        a = "from itertools import filterfalse, foo"
-        self.check(b, a)
-
-        b = "from itertools import bar, ifilterfalse, foo"
-        a = "from itertools import bar, filterfalse, foo"
-        self.check(b, a)
+    def test_ifilter_and_zip_longest(self):
+        for name in "filterfalse", "zip_longest":
+            b = "from itertools import i%s" % (name,)
+            a = "from itertools import %s" % (name,)
+            self.check(b, a)
+
+            b = "from itertools import imap, i%s, foo" % (name,)
+            a = "from itertools import %s, foo" % (name,)
+            self.check(b, a)
+
+            b = "from itertools import bar, i%s, foo" % (name,)
+            a = "from itertools import bar, %s, foo" % (name,)
+            self.check(b, a)
 
     def test_import_star(self):
         s = "from itertools import *"

-- 
Repository URL: http://hg.python.org/cpython