Automatically cleanup super calls for python3

super.py

import pathlib
import ast
import sys

import asttokens
import black


def walk(parent):
    for path in parent.iterdir():
        if path.is_dir():
            yield from walk(path)
        elif path.suffix == ".py":
            yield path


def get_replacements(tree):
    visitor = Visitor(tree)
    visitor.visit(tree.tree)
    return visitor.replacements


def replace(src, replacements):
    chunks = []
    end = len(src)
    for (start, stop), mod in reversed(replacements):
        chunks.append(src[stop:end])
        chunks.append(mod)
        end = start
    chunks.append(src[0:end])

    return "".join(reversed(chunks))


class Visitor(ast.NodeVisitor):
    def __init__(self, tree):
        self.tree = tree
        self.replacements = []
        self.stack = []

    def visit_ClassDef(self, node):
        self.stack.append(node.name)
        self.generic_visit(node)
        self.stack.pop()

    def visit_Call(self, node):
        self.generic_visit(node)
        if not isinstance(node.func, ast.Name):
            return
        if node.func.id != "super":
            return
        if not node.args:
            return
        if node.args[0].id != self.stack[-1]:
            print(
                f"Ignoring super call to {node.args[0].id} "
                f"while in {self.stack[-1]} context"
            )
            return
        if node.args[1].id not in ("self", "cls"):
            print("Ignoring super call not to self")
            return
        src = "super()"
        self.replacements.append((self.tree.get_text_range(node), src))


folder = sys.argv[1] if len(sys.argv) > 1 else "."

for path in walk(pathlib.Path(folder)):
    src = path.read_text()

    try:
        tree = asttokens.ASTTokens(src, filename=path, parse=True)
    except SyntaxError:
        print(f"Cannot parse {path}")
        continue

    print(f"  {path}")
    replacements = get_replacements(tree)
    if not replacements:
        continue

    print(f"Modifying {len(replacements)} calls in {path}")
    src = replace(src, replacements)

    src = black.format_str(src, line_length=79)
    path.write_text(src)