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 ESC[4;38;5;81mFixItertoolsImports(ESC[4;38;5;149mfixer_baseESC[4;38;5;149m.ESC[4;38;5;149mBaseFix):
       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 ('imap', 'izip', 'ifilter'):
      32                  child.value = None
      33                  child.remove()
      34              elif member_name in ('ifilterfalse', 'izip_longest'):
      35                  node.changed()
      36                  name_node.value = ('filterfalse' if member_name[1] == 'f'
      37                                     else 'zip_longest')
      38  
      39          # Make sure the import statement is still sane
      40          children = imports.children[:] or [imports]
      41          remove_comma = True
      42          for child in children:
      43              if remove_comma and child.type == token.COMMA:
      44                  child.remove()
      45              else:
      46                  remove_comma ^= True
      47  
      48          while children and children[-1].type == token.COMMA:
      49              children.pop().remove()
      50  
      51          # If there are no imports left, just get rid of the entire statement
      52          if (not (imports.children or getattr(imports, 'value', None)) or
      53              imports.parent is None):
      54              p = node.prefix
      55              node = BlankLine()
      56              node.prefix = p
      57              return node