pip-dump

#!/usr/bin/env python
from __future__ import absolute_import
import os.path
import glob
import argparse
import logging
import sys
import tempfile
from itertools import dropwhile, takewhile, chain
from functools import partial
from subprocess import check_call as _check_call

try:
    from subprocess import check_output as _check_output
except ImportError:
    import subprocess

    def _check_output(*args, **kwargs):
        process = subprocess.Popen(stdout=subprocess.PIPE, *args, **kwargs)
        output, _ = process.communicate()
        retcode = process.poll()
        if retcode:
            error = subprocess.CalledProcessError(retcode, args[0])
            error.output = output
            raise error
        return output

check_call = partial(_check_call, shell=True)
check_output = partial(_check_output, shell=True)

# Constants
DEFAULT_REQUIREMENTS_FILE = u'requirements.txt'
GLOB_PATTERNS = (u'*requirements.txt', u'requirements[_-]*.txt', u'requirements/*.txt')
PIP_IGNORE_FILE = u'.pipignore'
SPLIT_PATTERN = u'## The following requirements were added by pip --freeze:'


class StdOutFilter(logging.Filter):
    def filter(self, record):
        return record.levelno in [logging.DEBUG, logging.INFO]


def setup_logging(verbose):
    if verbose:
        level = logging.DEBUG
    else:
        level = logging.INFO

    format = '%(message)s'

    logger = logging.getLogger(u'pip-dump')

    stdout_handler = logging.StreamHandler(sys.stdout)
    stdout_handler.addFilter(StdOutFilter())
    stdout_handler.setFormatter(logging.Formatter(format))
    stdout_handler.setLevel(logging.DEBUG)

    stderr_handler = logging.StreamHandler(sys.stderr)
    stderr_handler.setFormatter(logging.Formatter(format))
    stderr_handler.setLevel(logging.WARNING)

    logger.setLevel(level)
    logger.addHandler(stderr_handler)
    logger.addHandler(stdout_handler)
    return logger


def parse_args():
    parser = argparse.ArgumentParser(
        description='Rewrites requirements.txt to match your virtualenv.'
    )
    parser.add_argument(
        '--verbose', '-v', action='store_true', default=False, help='Show more output'
    )
    parser.add_argument(
        '--no-sort', action='store_false', default=True, help="Don't sort package names (default is to sort)", dest="should_sort"
    )
    parser.add_argument('files', nargs='*')
    return parser.parse_args()


def pip_partition(lines):
    no_split_match = lambda line: line != SPLIT_PATTERN
    split_match = lambda line: line == SPLIT_PATTERN
    first = takewhile(no_split_match, lines)
    second = dropwhile(split_match, dropwhile(no_split_match, lines))
    return list(first), list(second)


def pip_info(filename):
    if not os.path.exists(filename):
        check_call('touch {0}'.format(filename))
    cmd = 'pip freeze -lr {0}'.format(filename,)
    raw = check_output(cmd)
    lines = raw.decode('utf-8').split('\n')
    return pip_partition(lines)


def append_lines(lines, filename):
    with open(filename, 'a') as f:
        for line in lines:
            f.write('{0}\n'.format(line))


def rewrite(filename, lines, should_sort=True):
    if should_sort:
        lines.sort(key=lambda s: s.lower())
    with open(filename, 'w') as f:
        for line in lines:
            line = line.strip()
            if line:
                f.write('{0}\n'.format(line))


def dump_requirements(files, should_sort=True):
    if not files:
        raise ValueError(u'Expected a list of at least one file name.')

    _, tmpfile = tempfile.mkstemp()
    existing_files = list(filter(os.path.exists, files))
    if existing_files:
        if should_sort:
            check_call(u'cat {0} | sort -fu > {1}'.format(' '.join(existing_files), tmpfile))
        else:
            check_call(u'cat {0} > {1}'.format(' '.join(existing_files), tmpfile))
    _, new = pip_info(tmpfile)
    check_call(u'rm {0}'.format(tmpfile))
    append_lines(new, files[0])

    for filename in files:
        if os.path.basename(filename) == PIP_IGNORE_FILE:  # never rewrite the pip ignore file
            continue
        pkgs, _ = pip_info(filename)
        rewrite(filename, pkgs, should_sort=should_sort)


def find_default_files():
    req_files = list(chain.from_iterable(glob.glob(pattern) for pattern in GLOB_PATTERNS))
    try:
        req_files.remove(DEFAULT_REQUIREMENTS_FILE)
    except ValueError:
        pass
    else:
        # Move the "default" requirements file to the front, so any new
        # packages will be added there
        req_files = [DEFAULT_REQUIREMENTS_FILE] + req_files

    # In case of no found files, use the default file
    if not req_files:
        req_files = [DEFAULT_REQUIREMENTS_FILE]

    if os.path.exists(PIP_IGNORE_FILE) and not PIP_IGNORE_FILE in req_files:
        req_files.append(PIP_IGNORE_FILE)
    return req_files


def main():
    args = parse_args()
    logger = setup_logging(args.verbose)

    if not args.files:
        args.files = find_default_files()
    dump_requirements(args.files, should_sort=args.should_sort)


if __name__ == '__main__':
    main()