condaenv

#!/usr/bin/env python2
import json
import os.path
import tempfile
import subprocess



def extract_conda(dat):
    deps = []
    for d in dat['dependencies']:
        try:
            d['pip']
            continue
        except TypeError:
            deps.append(d)
    return deps

def extract_pip(dat):
    for d in dat['dependencies']:
        try:
            return d['pip']
        except TypeError:
            continue
    return []

def load_yaml(path):
    try:
        import yaml
    except ImportError:
        raise SystemExit('ERROR: missing yaml - please run: conda install -n root pyyaml')
    with open(path, 'rt') as f:
        return yaml.load(f, Loader=yaml.Loader)


def get_python_interpreter(conda, env_name):
    import json

    try:
        dat = json.loads(subprocess.check_output([conda, 'env', 'list', '--json']))
        for e in dat['envs']:
            if os.path.basename(e) == env_name:
                return os.path.join(e, 'bin', 'python')
    except:
        pass

    return None


def conda_version(conda_bin):
    proc = subprocess.Popen([conda_bin, '--version'],
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)
    stdout, stderr = proc.communicate()
    assert proc.returncode == 0
    assert stderr.startswith('conda ')
    version = stderr.strip()[6:].split('.')
    return tuple(int(x) if x.isdigit() else x for x in version)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('-f','--file')
    parser.add_argument('--version', action='store_true')
    parser.add_argument('--force', action='store_true')
    parser.add_argument('-n','--name')
    parser.add_argument('--download-only', action='store_true')
    parser.add_argument('--download-dir', default='./CACHE', help='used when --download-only requested default="./CACHE"')
    parser.add_argument('-y', '--yes', action='store_true')
    parser.add_argument('--update', action='store_true', help='update instead of creating the env')
    parser.add_argument('--conda-binary', default='conda', help='path to conda binary')
    args = parser.parse_args()

    if args.version:
        print "condaenv 0.1.3"
        parser.exit(0)

    if not args.file:
        parser.error('argument -f/--file is required')

    dat = load_yaml(args.file)

    name = args.name
    try:
        if not name:
            name = dat['name']
    except KeyError:
        pass
    if not name:
        parser.error('-n, --name=NAME is required')

    if not args.update:
        proc_cmd = 'create'
    else:
        proc_cmd = 'install'

    conda_bin = args.conda_binary

    # check conda is installed
    conda_ver = conda_version(conda_bin)
    print "conda version: {}".format(conda_ver)

    conda_root = json.loads(subprocess.check_output(['conda', 'info', '--json']))['root_prefix']

    proc_args = [conda_bin, proc_cmd, '-n', name]
    proc_env = os.environ.copy()  # shell environment to run conda in

    if args.download_only:
        proc_args.append('--download-only')
        pkgs_dirs = os.path.abspath(os.path.expanduser(args.download_dir))
        if not os.path.exists(pkgs_dirs):
            os.makedirs(pkgs_dirs)
        elif not os.path.isdir(pkgs_dirs):
            raise SystemExit('ERROR: requested download dir "%s" not a directory' % pkgs_dirs)
        proc_env['CONDA_PKGS_DIRS'] = pkgs_dirs

    if args.force and not args.update:
        # does not error if env doesnt exist
        # fixme: this should be different depending on if we are pretending
        # to be conda env create or conda env upgrade
        subprocess.check_call([conda_bin, 'remove', '-n', name, '--all', '-y'])

    if args.yes:
        proc_args.append('-y')

    cdeps = extract_conda(dat)

    chans = {}
    for chan in dat.get('channels', []):
        proc_args.extend(['--channel', chan])
        chans[chan] = True

    for d in cdeps:
        try:
            chan, _ = d.split('::')
        except ValueError:
            parser.error('dependencies must indicate channel: %s' % d)

        if chan not in chans:
            proc_args.extend(['--channel', chan])
            chans[chan] = True

    # this command should be the whole world. user defaults have no place here
    proc_args.extend(['--override-channels', '--no-channel-priority'])

    # add dependencies to command line
    proc_args.extend(cdeps)

    if cdeps:
        print "====CONDAENV:", ' '.join(proc_args)
        try:
            subprocess.check_call(proc_args, env=proc_env)
        except subprocess.CalledProcessError:
            print "\nIF CONDA HAS AN UNFALSIFIABLE ERROR, IGNORE ITS BAD ADVICE"
            print "YOU CAN TYPE FOR EXAMPLE TO LEARN ABOUT PACKAGE DEPENDENCIES"
            print "\n\tconda search --info 'conda-forge::scikit-learn=0.19.0'\n"
            parser.exit(1)
    else:
        print "====CONDAENV: no conda dependencies"

    # command exited successfully

    pyexe = get_python_interpreter(conda_bin, name)
    print "====CONDAENV: interpreter=%s" % pyexe

    if pyexe:
        # we succeeded / user didnt abort
        pdeps = extract_pip(dat)
        if pdeps:
            fd, path = tempfile.mkstemp(suffix='requirements.txt', text=True)
            with os.fdopen(fd, 'wt') as f:
                f.write('\n'.join(pdeps))

            proc_args = [pyexe, '-m', 'pip', 'install', '--no-cache-dir', '-r', path]

            # we need to do this in a full subshell so that all the PATH and friends are
            # available for packages that build native code and they themselves call out
            # to pkg-config subshells
            if conda_ver >= (4, 5):
                conda_activate = [
                    'source', os.path.join(conda_root, 'etc', 'profile.d', 'conda.sh'), '&&',
                    'conda', 'activate', name
                ]
            else:
                conda_activate = [
                    'source', 'activate', name
                ]
            proc_args = ' '.join(conda_activate + ['&&'] + proc_args)
            shell = True
            executable = '/bin/bash'

            print "====CONDAENV:", proc_args
            subprocess.check_call(proc_args, shell=shell, executable=executable)

#             os.unlink(path)