Requirements solver

requirements.py

# -*- coding: utf-8 -*-
"""
Generates a requirements.txt from any packages
in the current working directory.
"""

import os
import glob
import contextlib
import functools
import argparse
import setuptools
import pkg_resources


class UnspecifiedRequirementError(AssertionError):
    pass


def _assert_requirement_pinned(requirement):
    if not requirement.specifier:
        raise UnspecifiedRequirementError(
            'Missing specifier for {!r}'.format(requirement.name)
        )
    return True


def _filter_requirements(
        predicate, output, install_requires=None,
        extras_require=None, dependency_links=None,
        *_args, **_kwargs
):
    if install_requires:
        output.update(
            str(requirement) for requirement in
            pkg_resources.parse_requirements(install_requires)
            if predicate(requirement, extra=None)
        )

    if dependency_links:
        output.update(dependency_links)

    if extras_require:
        for extra, requirements in extras_require.items():
            output.update(
                str(requirement) for requirement in
                pkg_resources.parse_requirements(requirements)
                if predicate(requirement, extra=extra)
            )


@contextlib.contextmanager
def patch(target, attr, fn):
    """
    Replace a function for the duration of the context manager
    """
    original = getattr(target, attr)
    try:
        setattr(target, attr, fn)
        yield
    finally:
        setattr(target, attr, original)


def find_requirements(predicate, *paths):
    """
    Return the list of requirements for a setup.py
    """
    requirements = set()

    setup = functools.partial(
        _filter_requirements, predicate, requirements
    )

    try:
        with patch(setuptools, 'setup', setup):
            for path in paths:
                with open(path, 'r') as handle:
                    exec(
                        handle.read(),
                        {'__file__': os.path.realpath(path)},
                        None
                    )
    except UnspecifiedRequirementError as err:
        raise UnspecifiedRequirementError(
            '{} in {}'.format(
                err.message, path
            )
        )

    for requirement in requirements:
        yield requirement


def main():
    """
    Print a requirements.txt based on the
    install_requires and dependency_links of our setup.py files
    """

    parser = argparse.ArgumentParser(
        description='Generate a requirements.txt from setup.pys'
    )
    parser.add_argument(
        '--extras',
        help='Generate an requirements.txt that only includes extra',
        action='append',
        default=[None]
    )
    parser.add_argument(
        'setups',
        nargs='+',
        help='Target setup.py files to scan for requirements'
    )
    args = parser.parse_args()

    print("""#
# autogenerated by requirements.py for build/osirium-venv
#
""")
    requirements = sorted(
        find_requirements(
            lambda requirement, extra: _assert_requirement_pinned(requirement) and extra in args.extras,
            *args.setups
        ),
        key=str.lower
    )

    for requirement in requirements:
        print(requirement)


if __name__ == '__main__':
    main()