1 """ Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """ 2 3 # Local imports 4 from lib2to3 import fixer_base 5 from lib2to3.fixer_util import BlankLine, syms, token 6 7 8 class 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
