aboutsummaryrefslogtreecommitdiff
path: root/delivery_merge/conda.py
blob: cc265f11940e3d965133883f852f08d267946f8a (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import os
import requests
import sys
from .utils import getenv, sh
from contextlib import contextmanager
from subprocess import run


ENV_ORIG = os.environ.copy()


class BadPlatform(Exception):
    pass


def conda_installer(ver, prefix='./miniconda3'):
    """ Install miniconda into a user-defined prefix and return its path

    :param ver: str: miniconda version (not conda version)
    :param prefix: str: path to install miniconda into
    :returns: str: absolute path to installation prefix
    :raises delivery_merge.conda.BadPlatform: when platform check fails
    :raises subprocess.CalledProcessError: via check_returncode method
    """
    assert isinstance(ver, str)
    assert isinstance(prefix, str)

    prefix = os.path.abspath(prefix)

    # Is miniconda already installed?
    if os.path.exists(prefix):
        print(f'{prefix}: exists', file=sys.stderr)
        return prefix

    name = 'Miniconda3'
    version = ver
    arch = 'x86_64'
    platform = sys.platform

    # Emit their installer's concept of "platform"
    if sys.platform == 'darwin':
        platform = 'MacOSX'
    elif sys.platform == 'linux':
        platform = 'Linux'
    else:
        raise BadPlatform(f'{sys.platform} is not supported.')

    url_root = 'https://repo.continuum.io/miniconda'
    installer = f'{name}-{version}-{platform}-{arch}.sh'
    url = f'{url_root}/{installer}'
    install_command = f'./{installer} -b -p {prefix}'.split()

    # Download installer
    if not os.path.exists(installer):
        with requests.get(url, stream=True) as data:
            with open(installer, 'wb') as fd:
                for chunk in data.iter_content(chunk_size=16384):
                    fd.write(chunk)
        os.chmod(installer, 0o755)

    # Perform installation
    run(install_command).check_returncode()

    return prefix


def conda_init_path(prefix):
    """ Redefines $PATH so subsequent shell calls use the just-installed
    miniconda prefix. This function will not continue prepending to $PATH
    so it's safe to call more than once.

    :param prefix: str: path to miniconda installation
    :returns: None
    """
    if os.environ['PATH'] != ENV_ORIG['PATH']:
        os.environ['PATH'] = ENV_ORIG['PATH']
    os.environ['PATH'] = ':'.join([os.path.join(prefix, 'bin'),
                                   os.environ['PATH']])


def conda_activate(env_name):
    """ Activate a conda environment

    Assume: `conda_init_path` as been called beforehand
    Warning: Arbitrary code execution is possible here due to `shell` usage.

    :param env_name: str: conda environment to activate
    :returns: dict: new runtime environment
    :raises subprocess.CalledProcessError: via check_returncode method
    """
    proc = run(f". activate {env_name} && env",
               capture_output=True,
               shell=True)
    proc.check_returncode()
    return getenv(proc.stdout.decode()).copy()


@contextmanager
def conda_env_load(env_name):
    """ A simple wrapper for `conda_activate`
    The current runtime environment is replaced and restored

    >>> with conda_env_load('some_env') as _:
    >>>     # do something

    :param env_name: str: conda environment to activate
    :returns: None
    """
    last = os.environ.copy()
    os.environ = conda_activate(env_name)
    try:
        yield
    finally:
        os.environ = last.copy()


def conda(*args):
    """ Execute conda shell commands

    :returns: subprocess.CompletedProcess object
    """
    return sh('conda', *args)


def conda_cmd_channels(conda_channels, override=True):
    """ Generate conda command arguments for handling channels
    :param conda_channels: list: URI to channel
    :param override: bool: channel order is preserved when True
    """
    assert isinstance(conda_channels, list)
    channels_result = '--override-channels ' if override else ''
    for channel in conda_channels:
        channels_result += f'-c {channel} '

    return channels_result