API for customizing the conda base environment.

condacolab.py

@@ -15,16 +15,23 @@
 import sys
 import shutil
 from datetime import datetime, timedelta
+from distutils.spawn import find_executable
 from pathlib import Path
 from subprocess import check_output, run, PIPE, STDOUT
 from textwrap import dedent
-from typing import Dict, AnyStr
+from typing import Dict, AnyStr, Iterable
 from urllib.request import urlopen
-from distutils.spawn import find_executable
-from IPython.display import display
+from urllib.error import HTTPError
 
+from IPython.display import display
 from IPython import get_ipython
 
+try:
+ from ruamel.yaml import YAML
+ from ruamel.yaml.comments import CommentedMap
+except ImportError as e:
+ raise RuntimeError("Could not find ruamel.yaml, plese install using `!pip install ruamel.yaml`!") from e
+
 try:
  import ipywidgets as widgets
  HAS_IPYWIDGETS = True
@@ -43,6 +50,7 @@
  "Surbhi Sharma <[email protected]>"
 )
 
+yaml=YAML()
 
 PREFIX = "/opt/conda"
 
@@ -87,12 +95,115 @@ def _run_subprocess(command, logs_filename):
  assert (task.returncode == 0), f"💥💔💥 The installation failed! Logs are available at `{logs_file_path}/{logs_filename}`."
 
 
+def _update_environment(
+ prefix:os.PathLike = PREFIX,
+ environment_file: str = None,
+ python_version: str = None,
+ specs: Iterable[str] = (),
+ channels: Iterable[str] = (),
+ pip_args: Iterable[str] = (),
+ extra_conda_args: Iterable[str] = (),
+ conda_exe: str = "conda",
+ ):
+ """
+ Install the dependencies in conda base environment during
+ the condacolab installion.
+
+ Parameters
+ ----------
+ prefix
+ Target location for the installation.
+ environment_file
+ Path or URL of the environment.yaml file to use for
+ updating the conda base enviornment.
+ python_version
+ Python version to use in the conda base environment, eg. "3.9".
+ specs
+ List of additional specifications (packages) to install.
+ channels
+ Comma separated list of channels to use in the conda
+ base environment.
+ pip_args
+ List of additional packages to be installed using pip.
+ extra_conda_args
+ Any extra conda arguments to be used during the installation.
+ """
+ os.makedirs("/var/condacolab", exist_ok=True)
+ environment_file_path = "/var/condacolab/environment.yaml"
+
+ # When environment.yaml file is not provided.
+ if environment_file is None:
+ env_details = {}
+ if channels:
+ env_details["channels"] = channels
+ if specs:
+ env_details["dependencies"] = specs
+ if python_version:
+ env_details["dependencies"] += [f"python={python_version}"]
+ if pip_args:
+ pip_args_dict = {"pip": pip_args}
+ env_details["dependencies"].append(pip_args_dict)
+
+ with open(environment_file_path, 'w') as f:
+ yaml.indent(mapping=2, sequence=4, offset=2)
+ yaml.dump(env_details, f)
+ else:
+ # If URL is given for environment.yaml file
+ if environment_file.startswith(("http://", "https://")):
+ try:
+ with urlopen(environment_file) as response, open(environment_file_path, "wb") as out:
+ shutil.copyfileobj(response, out)
+ except HTTPError as e:
+ raise HTTPError("The URL you entered is not working, please check it again.") from e
+
+ # If path is given for environment.yaml file
+ else:
+ shutil.copy(environment_file, environment_file_path)
+
+ with open(environment_file_path, 'r') as f:
+ env_details = yaml.load(f.read())
+
+ for key in env_details:
+ if channels and key == "channels":
+ env_details["channels"].extend(channels)
+ if key == "dependencies":
+ if specs:
+ env_details["dependencies"].extend(specs)
+ if python_version:
+ env_details["dependencies"].extend([f"python={python_version}"])
+ if pip_args:
+ for element in env_details["dependencies"]:
+ # if pip dependencies are already specified.
+ if isinstance(element, CommentedMap) and "pip" in element:
+ element["pip"].extend(pip_args)
+ break
+ # if there are no pip dependencies specified in the yaml file.
+ else:
+ pip_args_dict = CommentedMap([("pip", [*pip_args])])
+ env_details["dependencies"].append(pip_args_dict)
+
+ with open(environment_file_path, 'w') as f:
+ f.truncate(0)
+ yaml.dump(env_details, f)
+
+ _run_subprocess(
+ [conda_exe, "env", "update", "-n", "base", "-f", environment_file_path, *extra_conda_args],
+ "environment_file_update.log",
+ )
+
+
 def install_from_url(
  installer_url: AnyStr,
  prefix: os.PathLike = PREFIX,
  env: Dict[AnyStr, AnyStr] = None,
  run_checks: bool = True,
  restart_kernel: bool = True,
+ environment_file: str = None,
+ python_version: str = None,
+ specs: Iterable[str] = (),
+ channels: Iterable[str] = (),
+ pip_args: Iterable[str] = (),
+ extra_conda_args: Iterable[str] = (),
 ):
  """
  Download and run a constructor-like installer, patching
@@ -186,24 +297,38 @@ def install_from_url(
  "pip_task.log"
  )
 
+ print("📦 Updating enviornment using YAML file...")
+
+ _update_environment(
+ prefix=prefix,
+ environment_file=environment_file,
+ specs=specs,
+ channels=channels,
+ python_version=python_version,
+ pip_args=pip_args,
+ extra_conda_args=extra_conda_args,
+ conda_exe=f"{prefix}/bin/{conda_exe}",
+ )
+
  env = env or {}
  bin_path = f"{prefix}/bin"
 
- os.rename(sys.executable, f"{sys.executable}.renamed_by_condacolab.bak")
- with open(sys.executable, "w") as f:
- f.write(
- dedent(
- f"""
- #!/bin/bash
- source {prefix}/etc/profile.d/conda.sh
- conda activate
- unset PYTHONPATH
- mv /usr/bin/lsb_release /usr/bin/lsb_release.renamed_by_condacolab.bak
- exec {bin_path}/python $@
- """
- ).lstrip()
- )
- run(["chmod", "+x", sys.executable])
+ if os.path.exists(sys.executable):
+ os.rename(sys.executable, f"{sys.executable}.renamed_by_condacolab.bak")
+ with open(sys.executable, "w") as f:
+ f.write(
+ dedent(
+ f"""
+ #!/bin/bash
+ source {prefix}/etc/profile.d/conda.sh
+ conda activate
+ unset PYTHONPATH
+ mv /usr/bin/lsb_release /usr/bin/lsb_release.renamed_by_condacolab.bak
+ exec {bin_path}/python $@
+ """
+ ).lstrip()
+ )
+ run(["chmod", "+x", sys.executable])
 
  taken = timedelta(seconds=round((datetime.now() - t0).total_seconds(), 0))
  print(f"⏲ Done in {taken}")
@@ -220,18 +345,26 @@ def install_from_url(
  else:
  print("🔁 Please restart kernel by clicking on Runtime > Restart runtime.")
 
+
 def install_mambaforge(
- prefix: os.PathLike = PREFIX, env: Dict[AnyStr, AnyStr] = None, run_checks: bool = True, restart_kernel: bool = True,
+ prefix: os.PathLike = PREFIX,
+ env: Dict[AnyStr, AnyStr] = None,
+ run_checks: bool = True,
+ restart_kernel: bool = True,
+ specs: Iterable[str] = (),
+ python_version: str = None,
+ channels: Iterable[str] = (),
+ environment_file: str = None,
+ extra_conda_args: Iterable[str] = (),
+ pip_args: Iterable[str] = (),
+
 ):
  """
  Install Mambaforge, built for Python 3.7.
-
  Mambaforge consists of a Miniconda-like distribution optimized
  and preconfigured for conda-forge packages, and includes ``mamba``,
  a faster ``conda`` implementation.
-
  Unlike the official Miniconda, this is built with the latest ``conda``.
-
  Parameters
  ----------
  prefix
@@ -244,7 +377,6 @@ def install_mambaforge(
  to add those yourself in the raw string. They will
  end up added to a line like ``exec env VAR=VALUE python3...``.
  For example, a value with spaces should be passed as::
-
  env={"VAR": '"a value with spaces"'}
  run_checks
  Run checks to see if installation was run previously.
@@ -256,15 +388,35 @@ def install_mambaforge(
  automatically and get a button instead to do it.
  """
  installer_url = r"https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh"
- install_from_url(installer_url, prefix=prefix, env=env, run_checks=run_checks, restart_kernel=restart_kernel)
-
+ install_from_url(
+ installer_url,
+ prefix=prefix, 
+ env=env,
+ run_checks=run_checks,
+ restart_kernel=restart_kernel,
+ specs=specs,
+ python_version=python_version,
+ channels=channels,
+ environment_file=environment_file,
+ extra_conda_args=extra_conda_args,
+ pip_args=pip_args,
+ )
 
 # Make mambaforge the default
 install = install_mambaforge
 
 
 def install_miniforge(
- prefix: os.PathLike = PREFIX, env: Dict[AnyStr, AnyStr] = None, run_checks: bool = True, restart_kernel: bool = True,
+ prefix: os.PathLike = PREFIX,
+ env: Dict[AnyStr, AnyStr] = None,
+ run_checks: bool = True,
+ restart_kernel: bool = True,
+ specs: Iterable[str] = (),
+ python_version: str = None,
+ channels: Iterable[str] = (),
+ environment_file: str = None,
+ extra_conda_args: Iterable[str] = (),
+ pip_args: Iterable[str] = (),
 ):
  """
  Install Mambaforge, built for Python 3.7.
@@ -298,11 +450,32 @@ def install_miniforge(
  automatically and get a button instead to do it.
  """
  installer_url = r"https://github.com/jaimergp/miniforge/releases/latest/download/Miniforge-colab-Linux-x86_64.sh"
- install_from_url(installer_url, prefix=prefix, env=env, run_checks=run_checks, restart_kernel=restart_kernel)
+ install_from_url(
+ installer_url,
+ prefix=prefix,
+ env=env,
+ run_checks=run_checks,
+ restart_kernel=restart_kernel,
+ specs=specs,
+ python_version=python_version,
+ channels=channels,
+ environment_file=environment_file,
+ extra_conda_args=extra_conda_args,
+ pip_args=pip_args,
+ )
 
 
 def install_miniconda(
- prefix: os.PathLike = PREFIX, env: Dict[AnyStr, AnyStr] = None, run_checks: bool = True, restart_kernel: bool = True,
+ prefix: os.PathLike = PREFIX,
+ env: Dict[AnyStr, AnyStr] = None,
+ run_checks: bool = True,
+ restart_kernel: bool = True,
+ specs: Iterable[str] = (),
+ python_version: str = None,
+ channels: Iterable[str] = (),
+ environment_file: str = None,
+ extra_conda_args: Iterable[str] = (),
+ pip_args: Iterable[str] = (),
 ):
  """
  Install Miniconda 4.12.0 for Python 3.7.
@@ -331,11 +504,32 @@ def install_miniconda(
  automatically and get a button instead to do it.
  """
  installer_url = r"https://repo.anaconda.com/miniconda/Miniconda3-py37_4.12.0-Linux-x86_64.sh"
- install_from_url(installer_url, prefix=prefix, env=env, run_checks=run_checks, restart_kernel=restart_kernel)
+ install_from_url(
+ installer_url,
+ prefix=prefix,
+ env=env,
+ run_checks=run_checks,
+ restart_kernel=restart_kernel,
+ specs=specs,
+ python_version=python_version,
+ channels=channels,
+ environment_file=environment_file,
+ extra_conda_args=extra_conda_args,
+ pip_args=pip_args,
+ )
 
 
 def install_anaconda(
- prefix: os.PathLike = PREFIX, env: Dict[AnyStr, AnyStr] = None, run_checks: bool = True, restart_kernel: bool = True,
+ prefix: os.PathLike = PREFIX,
+ env: Dict[AnyStr, AnyStr] = None,
+ run_checks: bool = True,
+ restart_kernel: bool = True,
+ specs: Iterable[str] = (),
+ python_version: str = None,
+ channels: Iterable[str] = (),
+ environment_file: str = None,
+ extra_conda_args: Iterable[str] = (),
+ pip_args: Iterable[str] = (),
 ):
  """
  Install Anaconda 2022.05, the latest version built
@@ -365,7 +559,19 @@ def install_anaconda(
  automatically and get a button instead to do it.
  """
  installer_url = r"https://repo.anaconda.com/archive/Anaconda3-2022.05-Linux-x86_64.sh"
- install_from_url(installer_url, prefix=prefix, env=env, run_checks=run_checks, restart_kernel=restart_kernel)
+ install_from_url( 
+ installer_url,
+ prefix=prefix,
+ env=env,
+ run_checks=run_checks,
+ restart_kernel=restart_kernel,
+ specs=specs,
+ python_version=python_version,
+ channels=channels,
+ environment_file=environment_file,
+ extra_conda_args=extra_conda_args,
+ pip_args=pip_args,
+ )
 
 
 def check(prefix: os.PathLike = PREFIX, verbose: bool = True):
@@ -382,10 +588,6 @@ def check(prefix: os.PathLike = PREFIX, verbose: bool = True):
  Print success message if True
  """
  assert find_executable("conda"), "💥💔💥 Conda not found!"
-
- pymaj, pymin = sys.version_info[:2]
- sitepackages = f"{prefix}/lib/python{pymaj}.{pymin}/site-packages"
- assert sitepackages in sys.path, f"💥💔💥 PYTHONPATH was not patched! Value: {sys.path}"
  assert all(
  not path.startswith("/usr/local/") for path in sys.path
  ), f"💥💔💥 PYTHONPATH include system locations: {[path for path in sys.path if path.startswith('/usr/local')]}!"