| """ Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """ |
|
|
| |
| from lib2to3 import fixer_base |
| from lib2to3.fixer_util import BlankLine, syms, token |
|
|
|
|
| class FixItertoolsImports(fixer_base.BaseFix): |
| BM_compatible = True |
| PATTERN = """ |
| import_from< 'from' 'itertools' 'import' imports=any > |
| """ %(locals()) |
|
|
| def transform(self, node, results): |
| imports = results['imports'] |
| if imports.type == syms.import_as_name or not imports.children: |
| children = [imports] |
| else: |
| children = imports.children |
| for child in children[::2]: |
| if child.type == token.NAME: |
| member = child.value |
| name_node = child |
| elif child.type == token.STAR: |
| |
| return |
| else: |
| assert child.type == syms.import_as_name |
| name_node = child.children[0] |
| member_name = name_node.value |
| if member_name in ('imap', 'izip', 'ifilter'): |
| child.value = None |
| child.remove() |
| elif member_name in ('ifilterfalse', 'izip_longest'): |
| node.changed() |
| name_node.value = ('filterfalse' if member_name[1] == 'f' |
| else 'zip_longest') |
|
|
| |
| children = imports.children[:] or [imports] |
| remove_comma = True |
| for child in children: |
| if remove_comma and child.type == token.COMMA: |
| child.remove() |
| else: |
| remove_comma ^= True |
|
|
| while children and children[-1].type == token.COMMA: |
| children.pop().remove() |
|
|
| |
| if (not (imports.children or getattr(imports, 'value', None)) or |
| imports.parent is None): |
| p = node.prefix |
| node = BlankLine() |
| node.prefix = p |
| return node |
|
|
