the-littlest-jupyterhub/tljh/conda.py

"""
Wrap conda commandline program
"""
import os
import subprocess
import json
import hashlib
import contextlib
import logging
import tempfile
import time

import requests

from tljh import utils
from tljh.utils import parse_version as V


def sha256_file(fname):
    """
    Return sha256 of a given filename

    Copied from https://stackoverflow.com/a/3431838
    """
    hash_sha256 = hashlib.sha256()
    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_sha256.update(chunk)
    return hash_sha256.hexdigest()


def check_miniconda_version(prefix, version):
    """
    Return true if a miniconda install with version exists at prefix
    """
    versions = get_mamba_versions(prefix)
    if "conda" not in versions:
        return False

    return V(versions["conda"]) >= V(version)


def get_mamba_versions(prefix):
    """Parse `mamba --version` output into a dict

    which looks like:

    mamba 1.1.0
    conda 22.11.1

    into:

    {
        "mamba": "1.1.0",
        "conda": "22.11.1",
    }
    """
    versions = {}
    try:
        out = (
            subprocess.check_output(
                [os.path.join(prefix, "bin", "mamba"), "--version"],
                stderr=subprocess.STDOUT,
            )
            .decode()
            .strip()
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        return versions
    for line in out.strip().splitlines():
        pkg, version = line.split()
        versions[pkg] = version
    return versions


@contextlib.contextmanager
def download_miniconda_installer(installer_url, sha256sum):
    """
    Context manager to download miniconda installer from a given URL

    This should be used as a contextmanager. It downloads miniconda installer
    of given version, verifies the sha256sum & provides path to it to the `with`
    block to run.
    """
    logger = logging.getLogger("tljh")
    logger.info(f"Downloading conda installer {installer_url}")
    with tempfile.NamedTemporaryFile("wb", suffix=".sh") as f:
        tic = time.perf_counter()
        r = requests.get(installer_url)
        r.raise_for_status()
        f.write(r.content)
        # Remain in the NamedTemporaryFile context, but flush changes, see:
        # https://docs.python.org/3/library/os.html#os.fsync
        f.flush()
        os.fsync(f.fileno())
        t = time.perf_counter() - tic
        logger.info(f"Downloaded conda installer {installer_url} in {t:.1f}s")

        if sha256_file(f.name) != sha256sum:
            raise Exception("sha256sum hash mismatch! Downloaded file corrupted")

        yield f.name


def fix_permissions(prefix):
    """Fix permissions in the install prefix

    For all files in the prefix, ensure that:
    - everything is owned by current user:group
    - nothing is world-writeable

    Run after each install command.
    """
    utils.run_subprocess(["chown", "-R", f"{os.getuid()}:{os.getgid()}", prefix])
    utils.run_subprocess(["chmod", "-R", "o-w", prefix])


def install_miniconda(installer_path, prefix):
    """
    Install miniconda with installer at installer_path under prefix
    """
    utils.run_subprocess(["/bin/bash", installer_path, "-u", "-b", "-p", prefix])
    # fix permissions on initial install
    # a few files have the wrong ownership and permissions initially
    # when the installer is run as root
    fix_permissions(prefix)


def ensure_conda_packages(prefix, packages):
    """
    Ensure packages (from conda-forge) are installed in the conda prefix.

    Note that conda seem to update dependencies by default, so there is probably
    no need to have a update parameter exposed for this function.
    """
    conda_executable = [os.path.join(prefix, "bin", "mamba")]
    abspath = os.path.abspath(prefix)
    # Let subprocess errors propagate
    # Explicitly do *not* capture stderr, since that's not always JSON!
    # Scripting conda is a PITA!
    # FIXME: raise different exception when using
    raw_output = subprocess.check_output(
        conda_executable
        + [
            "install",
            "-c",
            "conda-forge",  # Make customizable if we ever need to
            "--json",
            "--prefix",
            abspath,
        ]
        + packages
    ).decode()
    # `conda install` outputs JSON lines for fetch updates,
    # and a undelimited output at the end. There is no reasonable way to
    # parse this outside of this kludge.
    filtered_output = "\n".join(
        [
            l
            for l in raw_output.split("\n")
            # Sometimes the JSON messages start with a \x00. The lstrip removes these.
            # conda messages seem to randomly throw \x00 in places for no reason
            if not l.lstrip("\x00").startswith('{"fetch"')
        ]
    )
    output = json.loads(filtered_output.lstrip("\x00"))
    if "success" in output and output["success"] == True:
        return
    fix_permissions(prefix)


def ensure_pip_packages(prefix, packages, upgrade=False):
    """
    Ensure pip packages are installed in the given conda prefix.
    """
    abspath = os.path.abspath(prefix)
    pip_executable = [os.path.join(abspath, "bin", "python"), "-m", "pip"]
    pip_cmd = pip_executable + ["install"]
    if upgrade:
        pip_cmd.append("--upgrade")
    utils.run_subprocess(pip_cmd + packages)
    fix_permissions(prefix)


def ensure_pip_requirements(prefix, requirements_path, upgrade=False):
    """
    Ensure pip packages from given requirements_path are installed in given conda prefix.

    requirements_path can be a file or a URL.
    """
    abspath = os.path.abspath(prefix)
    pip_executable = [os.path.join(abspath, "bin", "python"), "-m", "pip"]
    pip_cmd = pip_executable + ["install"]
    if upgrade:
        pip_cmd.append("--upgrade")
    utils.run_subprocess(pip_cmd + ["--requirement", requirements_path])
    fix_permissions(prefix)