Lib/lib2to3/fixes/fix_itertools_imports.py

   1 """ Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """
   2
   3 # Local imports
   4 from lib2to3 import fixer_base
   5 from lib2to3.fixer_util import BlankLine, syms, token
   6
   7
   8 class FixItertoolsImports(fixer_base.BaseFix):
   9     BM_compatible = True
  10     PATTERN = """
  11               import_from< 'from' 'itertools' 'import' imports=any >
  12               """ %(locals())
  13
  14     def transform(self, node, results):
  15         imports = results['imports']
  16         if imports.type == syms.import_as_name or not imports.children:
  17             children = [imports]
  18         else:
  19             children = imports.children
  20         for child in children[::2]:
  21             if child.type == token.NAME:
  22                 member = child.value
  23                 name_node = child
  24             elif child.type == token.STAR:
  25                 # Just leave the import as is.
  26                 return
  27             else:
  28                 assert child.type == syms.import_as_name
  29                 name_node = child.children[0]
  30             member_name = name_node.value
  31             if member_name in (u'imap', u'izip', u'ifilter'):
  32                 child.value = None
  33                 child.remove()
  34             elif member_name == u'ifilterfalse':
  35                 node.changed()
  36                 name_node.value = u'filterfalse'
  37
  38         # Make sure the import statement is still sane
  39         children = imports.children[:] or [imports]
  40         remove_comma = True
  41         for child in children:
  42             if remove_comma and child.type == token.COMMA:
  43                 child.remove()
  44             else:
  45                 remove_comma ^= True
  46
  47         while children and children[-1].type == token.COMMA:
  48             children.pop().remove()
  49
  50         # If there are no imports left, just get rid of the entire statement
  51         if (not (imports.children or getattr(imports, 'value', None)) or
  52             imports.parent is None):
  53             p = node.prefix
  54             node = BlankLine()
  55             node.prefix = p
  56             return node