OILS / devtools / fixes / fix_itertools_imports.py View on Github | oilshell.org

136 lines, 91 significant
1""" Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """
2
3# Local imports
4from lib2to3 import fixer_base
5from lib2to3 import pytree
6from lib2to3.fixer_util import BlankLine, syms, token
7
8
9# Unused: from fixes/fix_import.py
10def 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
31class 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