def main()

in build_docker.py [0:0]

• 48 lines of code
• 9 McCabe index (conditional complexity)

def main():
    frameworks = {
        dirname for dirname in os.listdir("docker_images") if dirname != "common"
    }
    framework_choices = frameworks.copy()
    framework_choices.add("all")

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "framework",
        type=str,
        choices=framework_choices,
        help="Which framework image to build.",
    )
    parser.add_argument(
        "--out",
        type=str,
        help="Where to store the new tags",
    )
    parser.add_argument(
        "--gpu",
        action="store_true",
        help="Build the GPU version of the model",
    )
    args = parser.parse_args()

    branch = (
        subprocess.check_output(["git", "rev-parse", "--abbrev-ref", "HEAD"])
        .decode("utf-8")
        .strip()
    )
    if branch != "main":
        raise Exception(f"Go to branch `main` ({branch})")

    print("Pulling")
    subprocess.run(["git", "pull"])

    if args.framework == "all":
        outputs = []
        for framework in frameworks:
            tag = build(framework, args.gpu)
            outputs.append((framework, tag))

    else:
        tag = build(args.framework, args.gpu)
        outputs = [(args.framework, tag)]

    for (framework, tag) in outputs:
        compute = "GPU" if args.gpu else "CPU"
        name = f"{framework.upper()}_{compute}_TAG"
        print(name, tag)
        if args.out:
            with open(args.out, "w") as f:
                f.write(f"{name}={tag}\n")