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