python (3.12.0)
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 ESC[4;38;5;81mFixItertoolsImports(ESC[4;38;5;149mfixer_baseESC[4;38;5;149m.ESC[4;38;5;149mBaseFix):
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 ('imap', 'izip', 'ifilter'):
32 child.value = None
33 child.remove()
34 elif member_name in ('ifilterfalse', 'izip_longest'):
35 node.changed()
36 name_node.value = ('filterfalse' if member_name[1] == 'f'
37 else '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