def transform(self, node, results): imports = results["imports"] if imports.type == syms.import_as_name or not imports.children: children = [imports] else: children = imports.children for child in children[::2]: if child.type == token.NAME: member = child.value name_node = child elif child.type == token.STAR: # Just leave the import as is. return else: assert child.type == syms.import_as_name name_node = child.children[0] member_name = name_node.value if member_name in ("imap", "izip", "ifilter"): child.value = None child.remove() elif member_name in ("ifilterfalse", "izip_longest"): node.changed() name_node.value = "filterfalse" if member_name[ 1] == "f" else "zip_longest" # Make sure the import statement is still sane children = imports.children[:] or [imports] remove_comma = True for child in children: if remove_comma and child.type == token.COMMA: child.remove() else: remove_comma ^= True while children and children[-1].type == token.COMMA: children.pop().remove() # If there are no imports left, just get rid of the entire statement if not (imports.children or getattr(imports, "value", None)) or imports.parent is None: p = node.prefix node = BlankLine() node.prefix = p return node
def transform_sort(self, node, results): sort_stmt = results["sort"] next_stmt = results["next"] list_call = results.get("list") simple_expr = results.get("expr") if list_call: list_call.replace(Name("sorted", prefix=list_call.prefix)) elif simple_expr: new = simple_expr.clone() new.prefix = "" simple_expr.replace(Call(Name("sorted"), [new], prefix=simple_expr.prefix)) else: raise RuntimeError("should not have reached here") sort_stmt.remove() btwn = sort_stmt.prefix # Keep any prefix lines between the sort_stmt and the list_call and # shove them right after the sorted() call. if "\n" in btwn: if next_stmt: # The new prefix should be everything from the sort_stmt's # prefix up to the last newline, then the old prefix after a new # line. prefix_lines = (btwn.rpartition("\n")[0], next_stmt[0].prefix) next_stmt[0].prefix = "\n".join(prefix_lines) else: assert list_call.parent assert list_call.next_sibling is None # Put a blank line after list_call and set its prefix. end_line = BlankLine() list_call.parent.append_child(end_line) assert list_call.next_sibling is end_line # The new prefix should be everything up to the first new line # of sort_stmt's prefix. end_line.prefix = btwn.rpartition("\n")[0]
def transform(self, node, results): new = BlankLine() new.prefix = node.prefix return new