WIP: initial mamba / micromamba support

.github/workflows/main.yml

@@ -77,7 +77,8 @@ jobs:
       fail-fast: false
       matrix:
         os: [macos-latest,ubuntu-latest,windows-latest]
-        pyver: ["3.8","3.10","3.12"]
+        pyver: ["3.9","3.12"]
+        solver: ["conda", "mamba", "micromamba"]
     steps:
     - name: Retrieve the source code
       uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4
@@ -114,6 +115,11 @@ jobs:
         [ 3.12 = ${{ matrix.pyver }} ] && export NBVER=7
         conda create -n testbase -c ./conda-bld nb_conda_kernels python=${{ matrix.pyver }} notebook=$NBVER pytest pytest-cov mock requests
         conda activate testbase
+        solver=${{ matrix.solver }}
+        [ $solver = conda ] || export CONDA_EXE=$(echo $CONDA_EXE | sed -E "s@^(.*/)conda(.*)@\\1${solver}\\2@")
+        [ $solver = mamba ] && export NBCK_NO_ACTIVATE_SCRIPT=yes
+        [ $solver = micromamba ] && export MAMBA_EXE=$CONDA_EXE
+        [ $solver = micromamba ] && export MAMBA_ROOT_PREFIX=$CONDA_ROOT
         python -m nb_conda_kernels list
         python -m pytest -v --cov=nb_conda_kernels tests 2>&1 | tee build.log
         # Because Windows refuses to preserve the error code

nb_conda_kernels/manager.py

@@ -10,19 +10,38 @@
 import psutil
 
 import os
-from os.path import join, split, dirname, basename, abspath
+from os.path import join, split, dirname, basename, abspath, exists
 from traitlets import Bool, Unicode, TraitError, validate
 
 from jupyter_client.kernelspec import KernelSpecManager, KernelSpec, NoSuchKernel
 
 CACHE_TIMEOUT = 60
 
-CONDA_EXE = os.environ.get("CONDA_EXE", "conda")
-
 RUNNER_COMMAND = ['python', '-m', 'nb_conda_kernels.runner']
 
 _canonical_paths = {}
 
+CONDA_EXE = None
+
+
+def _conda_exe():
+    global CONDA_EXE
+    if CONDA_EXE is not None:
+        return CONDA_EXE
+    for evar in ("CONDA_EXE", "MAMBA_EXE"):
+        CONDA_EXE = os.environ.get(evar)
+        if CONDA_EXE and exists(CONDA_EXE):
+            return CONDA_EXE
+    paths = os.environ.get("PATH").split(os.pathsep)
+    ext = ".exe" if sys.platform.startswith('win') else ""
+    for pname in ("conda", "mamba", "micromamba"):
+        for pdir in paths:
+            CONDA_EXE = join(pdir, pname + ext)
+            if exists(CONDA_EXE):
+                return CONDA_EXE
+    CONDA_EXE = ""
+    return CONDA_EXE
+
 
 def _canonicalize(path):
     """
@@ -123,9 +142,7 @@ def __init__(self, **kwargs):
         if not self._kernel_user:
             self._kernel_prefix = sys.prefix if self.kernelspec_path == "--sys-prefix" else self.kernelspec_path
 
-        self.log.info(
-            "nb_conda_kernels | enabled, %s kernels found.", len(self._conda_kspecs)
-        )
+        self.log.info("nb_conda_kernels | %d kernels found.", len(self._conda_kspecs))
 
     @staticmethod
     def clean_kernel_name(kname):
@@ -152,57 +169,91 @@ def _conda_info(self):
         """
 
         def get_conda_info_data():
-          # This is to make sure that subprocess can find 'conda' even if
-          # it is a Windows batch file---which is the case in non-root
-          # conda environments.
-          shell = CONDA_EXE == 'conda' and sys.platform.startswith('win')
-          try:
-            # Let json do the decoding for non-ASCII characters
-            out = subprocess.check_output([CONDA_EXE, "info", "--json"], shell=shell)
-            conda_info = json.loads(out)
-            return conda_info, None
-          except Exception as err:
-            return None, err
-          finally:
-             self.wait_for_child_processes_cleanup()
+            global CONDA_EXE
+            first_log = CONDA_EXE is None
+            conda_exe = _conda_exe()
+
+            if not first_log:
+                msg = None
+            elif conda_exe:
+                msg = "enabled: " + conda_exe
+            else:
+                msg = "could not find conda or mamba"
+            if not conda_exe:
+                return None, msg
+
+            try:
+               # Let json do the decoding for non-ASCII characters
+               out = subprocess.check_output([conda_exe, "info", "--json"])
+               conda_info = json.loads(out)
+               if 'envs' not in conda_info:
+                   # Micromamba does not include the envs list by default
+                   out = subprocess.check_output([conda_exe, "env", "list", "--json"])
+                   conda_info.update(json.loads(out))
+            except Exception as err:
+                msg = "error reading conda info: " + str(err)
+                return None, msg
+
+            finally:
+               self.wait_for_child_processes_cleanup()
+
+            # We moved the post-processing here so we can handle the conda/micromamba
+            # differences in one place
+            envs = list(map(_canonicalize, conda_info.get('envs') or ()))
+            base_prefix = _canonicalize(conda_info.get('conda_prefix') or conda_info.get('base environment'))
+            version = conda_info.get('conda_version') or conda_info.get('micromamba version')
+            if base_prefix not in envs:
+                # Older versions of conda do not include base_prefix in the env list
+                envs.insert(0, base_prefix)
+
+            return {
+                "conda_exe": conda_exe,
+                "conda_version": version,
+                "conda_prefix": base_prefix,
+                "envs": envs,
+            }, msg
 
         class CondaInfoThread(threading.Thread):
-          def run(self):
-            self.out, self.err = get_conda_info_data()
+            def run(self):
+              self.out, self.err = get_conda_info_data()
 
         expiry = self._conda_info_cache_expiry
         t = self._conda_info_cache_thread
 
         # cache is empty
+        msg, level = None, "debug"
         if expiry is None:
-          self.log.debug("nb_conda_kernels | refreshing conda info (blocking call)")
-          conda_info, err = get_conda_info_data()
-          if conda_info is None:
-            self.log.error("nb_conda_kernels | couldn't call conda:\n%s", err)
-          self._conda_info_cache = conda_info
-          self._conda_info_cache_expiry = time.time() + CACHE_TIMEOUT
+            conda_info, msg = get_conda_info_data()
+            if msg:
+                level = "info" if conda_info else "error"
+            else:
+                msg = "refreshing conda info (blocking call)"
+            self._conda_info_cache = conda_info
+            self._conda_info_cache_expiry = time.time() + CACHE_TIMEOUT
 
         # subprocess just finished
         elif t and not t.is_alive():
-          t.join()
-          conda_info = t.out
-          if conda_info is None:
-            self.log.error("nb_conda_kernels | couldn't call conda:\n%s", t.err)
-          else:
-            self.log.debug("nb_conda_kernels | collected conda info (async call)")
-          self._conda_info_cache = conda_info
-          self._conda_info_cache_expiry = time.time() + CACHE_TIMEOUT
-          self._conda_info_cache_thread = None
+            t.join()
+            conda_info, msg = t.out, t.err
+            if msg:
+                level = "info" if conda_info else "error"
+            else:
+                msg = "collected conda info (async call)"
+            self._conda_info_cache = conda_info
+            self._conda_info_cache_expiry = time.time() + CACHE_TIMEOUT
+            self._conda_info_cache_thread = None
 
         # cache expired
         elif not t and expiry < time.time():
-          self.log.debug("nb_conda_kernels | refreshing conda info (async call)")
-          t = CondaInfoThread()
-          t.start()
-          self._conda_info_cache_thread = t
+            msg = "refreshing conda info (async call)"
+            t = CondaInfoThread()
+            t.start()
+            self._conda_info_cache_thread = t
 
-        # else, just return cache
+        if msg:
+            getattr(self.log, level)("nb_conda_kernels | %s", msg)
 
+        # else, just return cache
         return self._conda_info_cache
 
     def _all_envs(self):
@@ -211,17 +262,10 @@ def _all_envs(self):
             canonical environment names as keys, and full paths as values.
         """
         conda_info = self._conda_info
-        envs = list(map(_canonicalize, conda_info['envs']))
-        base_prefix = _canonicalize(conda_info['conda_prefix'])
+        envs = conda_info['envs']
+        base_prefix = conda_info['conda_prefix']
         envs_prefix = join(base_prefix, 'envs')
         build_prefix = join(base_prefix, 'conda-bld', '')
-        # Older versions of conda do not seem to include the base prefix
-        # in the environment list, but we do want to scan that
-        if base_prefix not in envs:
-            envs.insert(0, base_prefix)
-        envs_dirs = conda_info['envs_dirs']
-        if not envs_dirs:
-            envs_dirs = [join(base_prefix, 'envs')]
         all_envs = {}
         for env_path in envs:
             if self.env_filter and self._env_filter_regex.search(env_path):
@@ -264,7 +308,9 @@ def _all_specs(self):
         all_specs = {}
         # We need to be able to find conda-run in the base conda environment
         # even if this package is not running there
-        conda_prefix = self._conda_info['conda_prefix']
+        conda_info = self._conda_info
+        conda_exe = conda_info['conda_exe']
+        conda_prefix = conda_info['conda_prefix']
         all_envs = self._all_envs()
         for env_name, env_path in all_envs.items():
             kspec_base = join(env_path, 'share', 'jupyter', 'kernels')
@@ -314,7 +360,7 @@ def _all_specs(self):
                     display_name += ' *'
                 spec['display_name'] = display_name
                 if env_path != sys.prefix:
-                    spec['argv'] = RUNNER_COMMAND + [conda_prefix, env_path] + spec['argv']
+                    spec['argv'] = RUNNER_COMMAND + [conda_exe, env_path] + spec['argv']
                 metadata = spec.get('metadata', {})
                 metadata.update({
                     'conda_env_name': env_name,

nb_conda_kernels/runner.py

@@ -1,6 +1,7 @@
 from __future__ import print_function
 
 import os
+from os.path import join, exists
 import sys
 import subprocess
 import locale
@@ -14,25 +15,39 @@ def exec_in_env(conda_prefix, env_path, *command):
     # Run the standard conda activation script, and print the
     # resulting environment variables to stdout for reading.
     is_current_env = env_path == sys.prefix
+    mamba_path = os.environ.get('MAMBA_EXE')
     if sys.platform.startswith('win'):
         if is_current_env:
             subprocess.Popen(list(command)).wait()
+            return
         else:
-            activate = os.path.join(conda_prefix, 'Scripts', 'activate.bat')
-            ecomm = [os.environ['COMSPEC'], '/S', '/U', '/C', '@echo', 'off', '&&',
-                    'chcp', '65001', '&&', 'call', activate, env_path, '&&',
-                    '@echo', 'CONDA_PREFIX=%CONDA_PREFIX%', '&&',] + list(command)
-            subprocess.Popen(ecomm).wait()
+            activate_path = join(conda_prefix, 'Scripts', 'activate.bat')
+            if exists(activate_path):
+                ecomm = [os.environ['COMSPEC'], '/S', '/U', '/C', '@echo', 'off', '&&',
+                        'chcp', '65001', '&&', 'call', activate_path, env_path, '&&',
+                        '@echo', 'CONDA_PREFIX=%CONDA_PREFIX%', '&&'] + list(command)
+                subprocess.Popen(ecomm).wait()
+            return
     else:
-        quoted_command = [quote(c) for c in command]
         if is_current_env:
-            os.execvp(quoted_command[0], quoted_command)
-        else:
-            activate = os.path.join(conda_prefix, 'bin', 'activate')
-            ecomm = ". '{}' '{}' && echo CONDA_PREFIX=$CONDA_PREFIX && exec {}".format(activate, env_path, ' '.join(quoted_command))
+            os.execvp(command[0], list(command))
+        activate = None
+        env_path = quote(env_path)
+        bin_path = join(conda_prefix, "bin")
+        if mamba_path and exists(mamba_path):
+            activate = 'eval "$({} shell activate {} --shell posix)"'.format(quote(mamba_path), env_path)
+        elif exists(join(bin_path, "activate")) and not os.environ.get('NBCK_NO_ACTIVATE_SCRIPT'):
+            activate = '. {}/activate {}'.format(quote(bin_path), env_path)
+        elif exists(join(bin_path, "conda")):
+            activate = 'eval "$({}/conda shell.posix activate {})"'.format(quote(bin_path), env_path)
+        if activate:
+            ecomm = (activate + '\nexec ') + ' '.join(quote(c) for c in command)
+            print(ecomm, file=sys.stderr)
             ecomm = ['sh' if 'bsd' in sys.platform else 'bash', '-c', ecomm]
             os.execvp(ecomm[0], ecomm)
 
+    raise RuntimeError('Could not determine an activation method')
+
 
 if __name__ == '__main__':
     exec_in_env(*(sys.argv[1:]))

testbed/build.sh

@@ -32,6 +32,7 @@ fi
 full_deactivate
 source $CONDA_ROOT/etc/profile.d/conda.sh
 conda activate base
+conda list
 
 if [ ! -f $CONDA_ROOT/.created ]; then
     conda config --prepend channels conda-forge --system

testbed/croot.yml

@@ -1,8 +1,10 @@
 name: conda
+channels:
+  - conda-forge
 dependencies:
-  - conda
-  - conda-build
-  - conda-verify
+  - conda >24.9
+  - mamba
+  - micromamba
   - notebook
   - jupyter_client
   - ipykernel

tests/test_config.py

@@ -37,8 +37,9 @@ def test_configuration():
         print('ERROR: Could not find conda find conda.')
         exit(-1)
     print(u'Current prefix: {}'.format(sys.prefix))
-    print(u'Root prefix: {}'.format(conda_info['root_prefix']))
-    print(u'Conda version: {}'.format(conda_info['conda_version']))
+    print(u'Conda prefix: {}'.format(conda_info['conda_prefix']))
+    conda_exe = os.path.basename(conda_info['conda_exe']).replace('.exe', '')
+    print(u'Conda version: {} {}'.format(conda_exe, conda_info['conda_version']))
     print(u'Environments:')
     for env in conda_info['envs']:
         print(u'  - {}'.format(env))
@@ -108,7 +109,8 @@ def test_kernel_name_format(monkeypatch, tmp_path, name_format, expected):
         "metadata": { "debugger": True }
     }
     mock_info = {
-        'conda_prefix': '/'
+        'conda_prefix': '/',
+        'conda_exe': '/conda.exe'
     }
     env_name = "dummy_env"
     def envs(*args):
@@ -214,7 +216,8 @@ def test_remove_kernelspec(tmp_path, kernel_name, expected):
 def test_kernel_metadata(monkeypatch, tmp_path, kernelspec):
 
     mock_info = {
-        'conda_prefix': '/'
+        'conda_prefix': '/',
+        'conda_exe': '/conda.exe'
     }
 
     def envs(*args):
@@ -258,6 +261,7 @@ def envs(*args):
 def test_kernel_metadata_debugger_override(monkeypatch, tmp_path, kernelspec):
 
     mock_info = {
+        'conda_exe': '/conda.exe',
         'conda_prefix': '/'
     }