[Python-checkins] r69673 - in sandbox/trunk/2to3/lib2to3: fixes/fix_itertools_imports.py tests/test_fixers.py

[benjamin.peterson]

Author: benjamin.peterson
Date: Mon Feb 16 16:38:22 2009
New Revision: 69673

Log:
fix handling of as imports #5279

Modified:
   sandbox/trunk/2to3/lib2to3/fixes/fix_itertools_imports.py
   sandbox/trunk/2to3/lib2to3/tests/test_fixers.py

Modified: sandbox/trunk/2to3/lib2to3/fixes/fix_itertools_imports.py
==============================================================================
--- sandbox/trunk/2to3/lib2to3/fixes/fix_itertools_imports.py	(original)
+++ sandbox/trunk/2to3/lib2to3/fixes/fix_itertools_imports.py	Mon Feb 16 16:38:22 2009
@@ -1,8 +1,9 @@
 """ Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """
 
 # Local imports
-from .. import fixer_base
-from ..fixer_util import BlankLine
+from lib2to3 import fixer_base
+from lib2to3.fixer_util import BlankLine, syms, token
+
 
 class FixItertoolsImports(fixer_base.BaseFix):
     PATTERN = """
@@ -11,34 +12,40 @@
 
     def transform(self, node, results):
         imports = results['imports']
-        children = imports.children[:] or [imports]
-        for child in children:
-            if not hasattr(child, 'value'):
-                # Handle 'import ... as ...'
-                continue
-            if child.value in ('imap', 'izip', 'ifilter'):
-                # The value must be set to none in case child == import,
-                # so that the test for empty imports will work out
+        if imports.type == syms.import_as_name or not imports.children:
+            children = [imports]
+        else:
+            children = imports.children
+        for child in children[::2]:
+            if child.type == token.NAME:
+                member = child.value
+                name_node = child
+            else:
+                assert child.type == syms.import_as_name
+                name_node = child.children[0]
+            member_name = name_node.value
+            if member_name in ('imap', 'izip', 'ifilter'):
                 child.value = None
                 child.remove()
-            elif child.value == 'ifilterfalse':
+            elif member_name == 'ifilterfalse':
                 node.changed()
-                child.value = 'filterfalse'
+                name_node.value = 'filterfalse'
 
         # Make sure the import statement is still sane
         children = imports.children[:] or [imports]
         remove_comma = True
         for child in children:
-            if remove_comma and getattr(child, 'value', None) == ',':
+            if remove_comma and child.type == token.COMMA:
                 child.remove()
             else:
                 remove_comma ^= True
 
-        if unicode(children[-1]) == ',':
+        if children[-1].type == token.COMMA:
             children[-1].remove()
 
         # If there are no imports left, just get rid of the entire statement
-        if not (imports.children or getattr(imports, 'value', None)):
+        if not (imports.children or getattr(imports, 'value', None)) or \
+                imports.parent is None:
             p = node.get_prefix()
             node = BlankLine()
             node.prefix = p

Modified: sandbox/trunk/2to3/lib2to3/tests/test_fixers.py
==============================================================================
--- sandbox/trunk/2to3/lib2to3/tests/test_fixers.py	(original)
+++ sandbox/trunk/2to3/lib2to3/tests/test_fixers.py	Mon Feb 16 16:38:22 2009
@@ -3402,6 +3402,18 @@
         a = "from itertools import bar as bang"
         self.check(b, a)
 
+        b = "from itertools import izip as _zip, imap, bar"
+        a = "from itertools import bar"
+        self.check(b, a)
+
+        b = "from itertools import imap as _map"
+        a = ""
+        self.check(b, a)
+
+        b = "from itertools import imap as _map, izip as _zip"
+        a = ""
+        self.check(b, a)
+
         s = "from itertools import bar as bang"
         self.unchanged(s)