Script to keep comments in requirements files after running pip-compile

pip-compile-keepcomments.py

#!/usr/bin/env python

import subprocess
import sys
import os.path

"""Script that runs pip-compile, but keeps comments.

Provide the same arguments as you would pip-compile, ensure that the filename is last.
The logic here is very simple and we don't care about edge cases, so it may break easily.
"""


def main():
    args = sys.argv[1:]
    input_filename = sys.argv[-1]
    filename_base, __ = os.path.splitext(input_filename)
    output_filename = filename_base + '.txt'

    # Read the existing output file and save for later
    if os.path.exists(output_filename):
        with open(output_filename) as f:
            existing_lines = [line for line in f if line.strip()]
    else:
        existing_lines = []

    # Run pip-compile with the same arguments as we got
    subprocess.run(['pip-compile', *args])

    if not os.path.exists(output_filename):
        return

    # For each line, if a longer version exists in the existing output file, use that.
    new_lines = []
    with open(output_filename) as f:
        for line in f:
            for existing_line in existing_lines:
                if line.strip() and existing_line.startswith(line.strip()):
                    new_lines.append(existing_line)
                    break
            else:
                new_lines.append(line)

    # Overwrite the output file with our new lines
    with open(output_filename, 'w') as f:
        for line in new_lines:
            f.write(line)


if __name__ == '__main__':
    main()