make_submission_zip.py

import argparse
from pathlib import Path
import zipfile


def parse_arguments():
    parser = argparse.ArgumentParser(
        description=(
            "Create a TacSIm benchmark submission ZIP "
            "from a directory of generated CSV files."
        )
    )
    parser.add_argument(
        "--input-dir",
        required=True,
    )
    parser.add_argument(
        "--output",
        default="submission.zip",
    )
    return parser.parse_args()


def main():
    args = parse_arguments()

    input_directory = Path(args.input_dir)
    csv_files = sorted(
        path
        for path in input_directory.rglob("*.csv")
        if path.is_file()
    )

    if not csv_files:
        raise ValueError(
            "No CSV files were found."
        )

    names = [path.name for path in csv_files]

    if len(names) != len(set(names)):
        raise ValueError(
            "CSV basenames must be unique."
        )

    output_path = Path(args.output)

    with zipfile.ZipFile(
        output_path,
        "w",
        zipfile.ZIP_DEFLATED,
    ) as archive:
        for csv_file in csv_files:
            archive.write(
                csv_file,
                csv_file.name,
            )

    print(
        f"Created {output_path} with "
        f"{len(csv_files)} CSV files."
    )


if __name__ == "__main__":
    main()