Updates PEP484 to modern types

update_pep484_types.py

import argparse
import ast


class PEP484ToModernTransformer(ast.NodeTransformer):
    GENERIC_MAP = {
        "Dict": "dict",
        "FrozenSet": "frozenset",
        "List": "list",
        "Set": "set",
        "Tuple": "tuple",
        "Type": "type",
    }

    def visit_Name(self, node: ast.Name) -> ast.AST:
        if node.id in self.GENERIC_MAP:
            return ast.Name(id=self.GENERIC_MAP[node.id], ctx=node.ctx)
        return self.generic_visit(node)

    def visit_Subscript(self, node: ast.Subscript) -> ast.AST:
        self.generic_visit(node)
        if not isinstance(node.value, ast.Name):
            return node
        type_name = node.value.id
        if type_name == "Optional":
            return ast.BinOp(
                left=node.slice, op=ast.BitOr(), right=ast.Constant(value=None)
            )
        elif type_name == "Union":
            types = (
                node.slice.elts if isinstance(node.slice, ast.Tuple) else [node.slice]
            )
            if not types:
                return node
            for t in types[1:]:
                result = ast.BinOp(left=types[0], op=ast.BitOr(), right=t)
            return result

        return node


def get_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("file", help="Python file to convert")

    return parser.parse_args()


if __name__ == "__main__":
    args = get_args()
    with open(args.file) as f:
        tree = ast.parse(f.read())
    new_tree = PEP484ToModernTransformer().visit(tree)
    ast.fix_missing_locations(new_tree)
    with open(args.file, "w") as f:
        f.write(ast.unparse(new_tree))
        f.write("\n")

    print(
        "Run `ruff check --fix` or similar tooling to remove unused imports and reformat"
    )