Quickly output requirements.txt of a Python project to the yaml format expected by cog

requirements_to_yaml.py

# usage:
# python -m venv venv
# source venv/bin/activate
# pip install -r requirements.txt
# pip freeze > requirements_frozen.txt
# python -m requirements_to_yaml --minimal

import re
import argparse

def main(requirements_file="requirements.txt", requirements_frozen_file="requirements_frozen.txt", minimal=False):
    """print pinned requirements in cog-compatible yaml form"""

    with open(requirements_file) as infile:
        reqs = infile.read().strip().split()

    constraint_pattern = r'[=><\!]'
    reqs_packages = [re.split(constraint_pattern, req)[0] for req in reqs]

    with open(requirements_frozen_file) as infile:
        reqs_frozen = infile.read().strip().split()

    if minimal:
        # if you only want the packages from requirements.txt, not the dependencies
        print("  python_packages:")
        print("\n".join([f'    - \"{req}\"' for req in reqs_frozen if re.split(constraint_pattern, req)[0] in reqs_packages]))
    else:
        print("  python_packages:")
        print("\n".join([f'    - \"{req}\"' for req in reqs_frozen]))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Print pinned requirements in cog-compatible yaml")
    parser.add_argument("-r","--requirements-file", default="requirements.txt", help="Path to requirements file")
    parser.add_argument("-f","--requirements-frozen-file", default="requirements_frozen.txt", help="Path to frozen requirements file")
    parser.add_argument("-m","--minimal", action="store_true", help="Print minimal requirements instead of requirements and their dependencies")

    args = parser.parse_args()

    main(requirements_file=args.requirements_file, requirements_frozen_file=args.requirements_frozen_file, minimal=args.minimal)