--- /dev/null
+""" Fixer for imports of itertools.(imap|ifilter|izip|ifilterfalse) """\r
+\r
+# Local imports\r
+from lib2to3 import fixer_base\r
+from lib2to3.fixer_util import BlankLine, syms, token\r
+\r
+\r
+class FixItertoolsImports(fixer_base.BaseFix):\r
+ BM_compatible = True\r
+ PATTERN = """\r
+ import_from< 'from' 'itertools' 'import' imports=any >\r
+ """ %(locals())\r
+\r
+ def transform(self, node, results):\r
+ imports = results['imports']\r
+ if imports.type == syms.import_as_name or not imports.children:\r
+ children = [imports]\r
+ else:\r
+ children = imports.children\r
+ for child in children[::2]:\r
+ if child.type == token.NAME:\r
+ member = child.value\r
+ name_node = child\r
+ elif child.type == token.STAR:\r
+ # Just leave the import as is.\r
+ return\r
+ else:\r
+ assert child.type == syms.import_as_name\r
+ name_node = child.children[0]\r
+ member_name = name_node.value\r
+ if member_name in (u'imap', u'izip', u'ifilter'):\r
+ child.value = None\r
+ child.remove()\r
+ elif member_name in (u'ifilterfalse', u'izip_longest'):\r
+ node.changed()\r
+ name_node.value = (u'filterfalse' if member_name[1] == u'f'\r
+ else u'zip_longest')\r
+\r
+ # Make sure the import statement is still sane\r
+ children = imports.children[:] or [imports]\r
+ remove_comma = True\r
+ for child in children:\r
+ if remove_comma and child.type == token.COMMA:\r
+ child.remove()\r
+ else:\r
+ remove_comma ^= True\r
+\r
+ while children and children[-1].type == token.COMMA:\r
+ children.pop().remove()\r
+\r
+ # If there are no imports left, just get rid of the entire statement\r
+ if (not (imports.children or getattr(imports, 'value', None)) or\r
+ imports.parent is None):\r
+ p = node.prefix\r
+ node = BlankLine()\r
+ node.prefix = p\r
+ return node\r