diff options
Diffstat (limited to 'Lib/lib2to3/fixes/fix_itertools_imports.py')
| -rw-r--r-- | Lib/lib2to3/fixes/fix_itertools_imports.py | 37 |
1 files changed, 22 insertions, 15 deletions
diff --git a/Lib/lib2to3/fixes/fix_itertools_imports.py b/Lib/lib2to3/fixes/fix_itertools_imports.py index ede8ad9..8a57f6d 100644 --- a/Lib/lib2to3/fixes/fix_itertools_imports.py +++ b/Lib/lib2to3/fixes/fix_itertools_imports.py @@ -1,8 +1,9 @@ """ Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """ # Local imports -from .. import fixer_base -from ..fixer_util import BlankLine +from lib2to3 import fixer_base +from lib2to3.fixer_util import BlankLine, syms, token + class FixItertoolsImports(fixer_base.BaseFix): PATTERN = """ @@ -11,34 +12,40 @@ class FixItertoolsImports(fixer_base.BaseFix): def transform(self, node, results): imports = results['imports'] - children = imports.children[:] or [imports] - for child in children: - if not hasattr(child, 'value'): - # Handle 'import ... as ...' - continue - if child.value in ('imap', 'izip', 'ifilter'): - # The value must be set to none in case child == import, - # so that the test for empty imports will work out + 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 + 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 child.value == 'ifilterfalse': + elif member_name == 'ifilterfalse': node.changed() - child.value = 'filterfalse' + name_node.value = 'filterfalse' # Make sure the import statement is still sane children = imports.children[:] or [imports] remove_comma = True for child in children: - if remove_comma and getattr(child, 'value', None) == ',': + if remove_comma and child.type == token.COMMA: child.remove() else: remove_comma ^= True - if str(children[-1]) == ',': + if children[-1].type == token.COMMA: children[-1].remove() # If there are no imports left, just get rid of the entire statement - if not (imports.children or getattr(imports, 'value', None)): + if not (imports.children or getattr(imports, 'value', None)) or \ + imports.parent is None: p = node.get_prefix() node = BlankLine() node.prefix = p |