Function main

scripts/prepare-for-update-assertions.py:108–129

(filename: str, function: str, dry_run: bool)

def main(filename: str, function: str, dry_run: bool) -> bool:
    with open(filename) as r:
        code = "".join(r.readlines())

    tree = cst.parse_module(code)
    transformer = SingleLineStatementTransformer(function)
    tree_updated = tree.visit(transformer)

    if dry_run:
        diff = "".join(difflib.unified_diff(code.splitlines(1), tree_updated.code.splitlines(1)))
        if diff:
            print(f"Diff of {filename}:")
            print(diff)
            print()
            return True
    else:
        if not tree_updated.deep_equals(tree):
            with open(filename, "w") as w:
                w.write(tree_updated.code)
            return True

    return False


def parse_args():