]>
Commit | Line | Data |
---|---|---|
4710c53d | 1 | """ Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """\r |
2 | \r | |
3 | # Local imports\r | |
4 | from lib2to3 import fixer_base\r | |
5 | from lib2to3.fixer_util import BlankLine, syms, token\r | |
6 | \r | |
7 | \r | |
8 | class FixItertoolsImports(fixer_base.BaseFix):\r | |
9 | BM_compatible = True\r | |
10 | PATTERN = """\r | |
11 | import_from< 'from' 'itertools' 'import' imports=any >\r | |
12 | """ %(locals())\r | |
13 | \r | |
14 | def transform(self, node, results):\r | |
15 | imports = results['imports']\r | |
16 | if imports.type == syms.import_as_name or not imports.children:\r | |
17 | children = [imports]\r | |
18 | else:\r | |
19 | children = imports.children\r | |
20 | for child in children[::2]:\r | |
21 | if child.type == token.NAME:\r | |
22 | member = child.value\r | |
23 | name_node = child\r | |
24 | elif child.type == token.STAR:\r | |
25 | # Just leave the import as is.\r | |
26 | return\r | |
27 | else:\r | |
28 | assert child.type == syms.import_as_name\r | |
29 | name_node = child.children[0]\r | |
30 | member_name = name_node.value\r | |
31 | if member_name in (u'imap', u'izip', u'ifilter'):\r | |
32 | child.value = None\r | |
33 | child.remove()\r | |
34 | elif member_name in (u'ifilterfalse', u'izip_longest'):\r | |
35 | node.changed()\r | |
36 | name_node.value = (u'filterfalse' if member_name[1] == u'f'\r | |
37 | else u'zip_longest')\r | |
38 | \r | |
39 | # Make sure the import statement is still sane\r | |
40 | children = imports.children[:] or [imports]\r | |
41 | remove_comma = True\r | |
42 | for child in children:\r | |
43 | if remove_comma and child.type == token.COMMA:\r | |
44 | child.remove()\r | |
45 | else:\r | |
46 | remove_comma ^= True\r | |
47 | \r | |
48 | while children and children[-1].type == token.COMMA:\r | |
49 | children.pop().remove()\r | |
50 | \r | |
51 | # If there are no imports left, just get rid of the entire statement\r | |
52 | if (not (imports.children or getattr(imports, 'value', None)) or\r | |
53 | imports.parent is None):\r | |
54 | p = node.prefix\r | |
55 | node = BlankLine()\r | |
56 | node.prefix = p\r | |
57 | return node\r |