improve runner

anaconda · Nov 6, 2024 · 90988c1 · 90988c1
1 parent b4a6121
commit 90988c1
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 14 deletions.
diff --git a/nb_conda_kernels/manager.py b/nb_conda_kernels/manager.py
@@ -308,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')
@@ -358,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,

diff --git a/nb_conda_kernels/runner.py b/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,24 +15,41 @@ 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]
+        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(cmd[0], command)
+        command = 'exec ' + ' '.join(command)
+        def _doit(activate):
+            print("|", activate, file=sys.stderr)
+            print("|", command, file=sys.stderr)
+            ecomm = activate + '\n' + command
             ecomm = ['sh' if 'bsd' in sys.platform else 'bash', '-c', ecomm]
             os.execvp(ecomm[0], ecomm)
+        env_path = quote(env_path)
+        if mamba_path and exists(mamba_path) and mamba_path.endswith('/micromamba'):
+            _doit('eval "$({} shell activate {} --shell posix)"'.format(quote(mamba_path), env_path))
+        bin_path = join(conda_prefix, "bin")
+        conda_path = join(bin_path, "conda")
+        if exists(conda_path):
+            _doit('eval "$({} shell.posix activate {})"'.format(quote(conda_path), env_path))
+        activate_path = join(bin_path, "activate")
+        if exists(activate_path):
+            _doit('. {} {}'.format(quote(activate_path), env_path))
+    raise RuntimeError('Could not determine an activation method')
 
 
 if __name__ == '__main__':

diff --git a/tests/test_config.py b/tests/test_config.py
@@ -109,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):
@@ -215,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):
@@ -259,6 +261,7 @@ def envs(*args):
 def test_kernel_metadata_debugger_override(monkeypatch, tmp_path, kernelspec):
 
     mock_info = {
+        'conda_exe': '/conda.exe',
         'conda_prefix': '/'
     }