1""" Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """
2
3# Local imports
4from lib2to3 import fixer_base
5from lib2to3.fixer_util import BlankLine, syms, token
6
7
8class 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 in (u'ifilterfalse', u'izip_longest'):
35                node.changed()
36                name_node.value = (u'filterfalse' if member_name[1] == u'f'
37                                   else u'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
58