1 """ Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """
4 from lib2to3 import fixer_base
5 from lib2to3.fixer_util import BlankLine, syms, token
8 class FixItertoolsImports(fixer_base.BaseFix):
11 import_from< 'from' 'itertools' 'import' imports=any >
14 def transform(self, node, results):
15 imports = results['imports']
16 if imports.type == syms.import_as_name or not imports.children:
19 children = imports.children
20 for child in children[::2]:
21 if child.type == token.NAME:
24 elif child.type == token.STAR:
25 # Just leave the import as is.
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'):
34 elif member_name == u'ifilterfalse':
36 name_node.value = u'filterfalse'
38 # Make sure the import statement is still sane
39 children = imports.children[:] or [imports]
41 for child in children:
42 if remove_comma and child.type == token.COMMA:
47 while children and children[-1].type == token.COMMA:
48 children.pop().remove()
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):