def main()

in scripts/rm_pt.py [0:0]

• 82 lines of code
• 1 McCabe index (conditional complexity)

def main():
    parser = argparse.ArgumentParser(
        description=(
            'Recursively delete checkpoint files from `root_dir`, '
            'but preserve checkpoint_best.pt and checkpoint_last.pt'
        )
    )
    parser.add_argument('root_dirs', nargs='*')
    parser.add_argument('--save-last', type=int, default=0, help='number of last checkpoints to save')
    parser.add_argument('--save-every', type=int, default=0, help='interval of checkpoints to save')
    parser.add_argument('--preserve-test', action='store_true',
                        help='preserve checkpoints in dirs that start with test_ prefix (default: delete them)')
    parser.add_argument('--delete-best', action='store_true', help='delete checkpoint_best.pt')
    parser.add_argument('--delete-last', action='store_true', help='delete checkpoint_last.pt')
    parser.add_argument('--no-dereference', action='store_true', help='don\'t dereference symlinks')
    args = parser.parse_args()

    files_to_desymlink = []
    files_to_preserve = []
    files_to_delete = []
    for root_dir in args.root_dirs:
        for root, _subdirs, files in os.walk(root_dir):
            if args.save_last > 0:
                to_save = last_n_checkpoints(files, args.save_last)
            else:
                to_save = []
            if args.save_every > 0:
                to_save += every_n_checkpoints(files, args.save_every)
            for file in files:
                if not pt_regexp.fullmatch(file):
                    continue
                full_path = os.path.join(root, file)
                if (
                    (
                        not os.path.basename(root).startswith('test_')
                        or args.preserve_test
                    )
                    and (
                        (file == 'checkpoint_last.pt' and not args.delete_last)
                        or (file == 'checkpoint_best.pt' and not args.delete_best)
                        or file in to_save
                    )
                ):
                    if os.path.islink(full_path) and not args.no_dereference:
                        files_to_desymlink.append(full_path)
                    else:
                        files_to_preserve.append(full_path)
                else:
                    files_to_delete.append(full_path)

    if len(files_to_desymlink) == 0 and len(files_to_delete) == 0:
        print('Nothing to do.')
        sys.exit(0)

    files_to_desymlink = sorted(files_to_desymlink)
    files_to_preserve = sorted(files_to_preserve)
    files_to_delete = sorted(files_to_delete)

    print('Operations to perform (in order):')
    if len(files_to_desymlink) > 0:
        for file in files_to_desymlink:
            print(' - preserve (and dereference symlink): ' + file)
    if len(files_to_preserve) > 0:
        for file in files_to_preserve:
            print(' - preserve: ' + file)
    if len(files_to_delete) > 0:
        for file in files_to_delete:
            print(' - delete: ' + file)
    while True:
        resp = input('Continue? (Y/N): ')
        if resp.strip().lower() == 'y':
            break
        elif resp.strip().lower() == 'n':
            sys.exit(0)

    print('Executing...')
    if len(files_to_desymlink) > 0:
        for file in files_to_desymlink:
            realpath = os.path.realpath(file)
            print('rm ' + file)
            os.remove(file)
            print('cp {} {}'.format(realpath, file))
            shutil.copyfile(realpath, file)
    if len(files_to_delete) > 0:
        for file in files_to_delete:
            print('rm ' + file)
            os.remove(file)