Remove Python 3 annotations from a codebase

remove_annotations.py

import functools
import pathlib
import sys

import libcst as cst


def main(root):
    root = pathlib.Path(root)
    fn = functools.partial(remove_annotations, annotations_remover=AnnotationRemover())
    try:
        import trailrunner

        trailrunner.walk_and_run([root], fn)
    except ModuleNotFoundError:
        for path in root.rglob("*") if root.is_dir() else [root]:
            fn(path)


class AnnotationRemover(cst.CSTTransformer):
    def on_visit(self, node: cst.CSTNode):
        return not isinstance(node, (cst.AnnAssign, cst.Annotation))

    def on_leave(self, original_node: cst.CSTNodeT, updated_node: cst.CSTNodeT):
        if isinstance(original_node, cst.Annotation):
            # FIXME: remove whitespace that is introduced by PEP8 around annotations inside signatures
            return cst.RemoveFromParent()
        elif isinstance(original_node, cst.AnnAssign):
            if original_node.value is None:
                return cst.RemoveFromParent()

            targets = [
                cst.AssignTarget(
                    target=original_node.target,
                    whitespace_before_equal=original_node.equal.whitespace_before,
                    whitespace_after_equal=original_node.equal.whitespace_after,
                )
            ]
            return cst.Assign(targets, original_node.value, original_node.semicolon)
        else:
            return updated_node


def remove_annotations(path, *, annotations_remover):
    if path.suffix != ".py":
        return

    with open(path) as file:
        with_annotations = cst.parse_module(file.read())

    without_annotations = with_annotations.visit(annotations_remover)

    if not without_annotations.deep_equals(with_annotations):
        with open(path, "w") as file:
            file.write(without_annotations.code)


if __name__ == "__main__":
    main(sys.argv[1])