| 1 | """ Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """
 | 
| 2 | 
 | 
| 3 | # Local imports
 | 
| 4 | from lib2to3 import fixer_base
 | 
| 5 | from lib2to3 import pytree
 | 
| 6 | from lib2to3.fixer_util import BlankLine, syms, token
 | 
| 7 | 
 | 
| 8 | 
 | 
| 9 | # Unused: from fixes/fix_import.py
 | 
| 10 | def traverse_imports(names):
 | 
| 11 |     """
 | 
| 12 |     Walks over all the names imported in a dotted_as_names node.
 | 
| 13 |     """
 | 
| 14 |     pending = [names]
 | 
| 15 |     while pending:
 | 
| 16 |         node = pending.pop()
 | 
| 17 |         if node.type == token.NAME:
 | 
| 18 |             yield node.value
 | 
| 19 |         elif node.type == syms.dotted_name:
 | 
| 20 |             yield "".join([ch.value for ch in node.children])
 | 
| 21 |         elif node.type == syms.dotted_as_name:
 | 
| 22 |             pending.append(node.children[0])
 | 
| 23 |         elif node.type == syms.dotted_as_names:
 | 
| 24 |             pending.extend(node.children[::-2])
 | 
| 25 |         else:
 | 
| 26 |             raise AssertionError("unknown node type")
 | 
| 27 | 
 | 
| 28 | 
 | 
| 29 | 
 | 
| 30 | 
 | 
| 31 | class FixItertoolsImports(fixer_base.BaseFix):
 | 
| 32 |     BM_compatible = True
 | 
| 33 |     PATTERN = """
 | 
| 34 |               import_from< 'from' imp=any 'import' ['('] imports=any [')'] >
 | 
| 35 |               """ %(locals())
 | 
| 36 | 
 | 
| 37 |     def transform(self, node, results):
 | 
| 38 |         #print('***')
 | 
| 39 | 
 | 
| 40 |         # lib2to3.pytree.Node
 | 
| 41 |         #print('NODE %s' % type(node))
 | 
| 42 | 
 | 
| 43 |         imp = results['imp']
 | 
| 44 |         if not isinstance(imp, pytree.Node):
 | 
| 45 |             # filter out from X import Y
 | 
| 46 |             return
 | 
| 47 | 
 | 
| 48 |         c0 = imp.children[0]
 | 
| 49 |         if c0.value != '_devbuild':
 | 
| 50 |             # Filter out
 | 
| 51 |             return
 | 
| 52 | 
 | 
| 53 |         imports = results['imports']
 | 
| 54 |         print()
 | 
| 55 |         print('I %r' % imports)
 | 
| 56 | 
 | 
| 57 |         n = len(imports.children)
 | 
| 58 |         to_remove = []
 | 
| 59 |         for i, child in enumerate(imports.children):
 | 
| 60 |             if child.value in ('value', 'value_e', 'value_t', 'value_str'):
 | 
| 61 |                 #print('NODE %r' % child)
 | 
| 62 |                 #print('NODE %r' % child.lineno)
 | 
| 63 |                 to_remove.append(child)
 | 
| 64 | 
 | 
| 65 |                 # Remove any preceding comma
 | 
| 66 |                 if i < n-1:
 | 
| 67 |                     after = imports.children[i+1]
 | 
| 68 |                     print('after %r' % after)
 | 
| 69 |                     if after.value == ',':
 | 
| 70 |                         to_remove.append(after)
 | 
| 71 | 
 | 
| 72 |         for n in to_remove:
 | 
| 73 |             imports.children.remove(n)
 | 
| 74 | 
 | 
| 75 |         # If there are no imports left, just get rid of the entire statement
 | 
| 76 |         # copied from fix_itertools_imports.py
 | 
| 77 |         if (not (imports.children or getattr(imports, 'value', None)) or
 | 
| 78 |             imports.parent is None):
 | 
| 79 |             p = node.prefix
 | 
| 80 |             node = BlankLine()
 | 
| 81 |             node.prefix = p
 | 
| 82 |             return node
 | 
| 83 | 
 | 
| 84 |         if to_remove:
 | 
| 85 |             return node
 | 
| 86 | 
 | 
| 87 |         return
 | 
| 88 | 
 | 
| 89 |         if imp == '_devbuild.gen.runtime_asdl':
 | 
| 90 |             print('IMPORT %s %s' % (imp, imports))
 | 
| 91 | 
 | 
| 92 |         #raise AssertionError()
 | 
| 93 |         return
 | 
| 94 | 
 | 
| 95 |         if imports.type == syms.import_as_name or not imports.children:
 | 
| 96 |             children = [imports]
 | 
| 97 |         else:
 | 
| 98 |             children = imports.children
 | 
| 99 |         for child in children[::2]:
 | 
| 100 |             if child.type == token.NAME:
 | 
| 101 |                 member = child.value
 | 
| 102 |                 name_node = child
 | 
| 103 |             elif child.type == token.STAR:
 | 
| 104 |                 # Just leave the import as is.
 | 
| 105 |                 return
 | 
| 106 |             else:
 | 
| 107 |                 assert child.type == syms.import_as_name
 | 
| 108 |                 name_node = child.children[0]
 | 
| 109 |             member_name = name_node.value
 | 
| 110 |             if member_name in ('imap', 'izip', 'ifilter'):
 | 
| 111 |                 child.value = None
 | 
| 112 |                 child.remove()
 | 
| 113 |             elif member_name in ('ifilterfalse', 'izip_longest'):
 | 
| 114 |                 node.changed()
 | 
| 115 |                 name_node.value = ('filterfalse' if member_name[1] == 'f'
 | 
| 116 |                                    else 'zip_longest')
 | 
| 117 | 
 | 
| 118 |         # Make sure the import statement is still sane
 | 
| 119 |         children = imports.children[:] or [imports]
 | 
| 120 |         remove_comma = True
 | 
| 121 |         for child in children:
 | 
| 122 |             if remove_comma and child.type == token.COMMA:
 | 
| 123 |                 child.remove()
 | 
| 124 |             else:
 | 
| 125 |                 remove_comma ^= True
 | 
| 126 | 
 | 
| 127 |         while children and children[-1].type == token.COMMA:
 | 
| 128 |             children.pop().remove()
 | 
| 129 | 
 | 
| 130 |         # If there are no imports left, just get rid of the entire statement
 | 
| 131 |         if (not (imports.children or getattr(imports, 'value', None)) or
 | 
| 132 |             imports.parent is None):
 | 
| 133 |             p = node.prefix
 | 
| 134 |             node = BlankLine()
 | 
| 135 |             node.prefix = p
 | 
| 136 |             return node
 |