Merge pull request #2207 from xhochy/mamba

asottile · web-flow · commit 12b482345b4e · 2022-01-15T17:57:13.000-05:00
Add mamba support to `language: conda`
diff --git a/pre_commit/languages/conda.py b/pre_commit/languages/conda.py
@@ -50,6 +50,15 @@ def in_env(
         yield
 
 
+def _conda_exe() -> str:
+    if os.environ.get('PRE_COMMIT_USE_MICROMAMBA'):
+        return 'micromamba'
+    elif os.environ.get('PRE_COMMIT_USE_MAMBA'):
+        return 'mamba'
+    else:
+        return 'conda'
+
+
 def install_environment(
         prefix: Prefix,
         version: str,
@@ -58,15 +67,17 @@ def install_environment(
     helpers.assert_version_default('conda', version)
     directory = helpers.environment_dir(ENVIRONMENT_DIR, version)
 
+    conda_exe = _conda_exe()
+
     env_dir = prefix.path(directory)
     with clean_path_on_failure(env_dir):
         cmd_output_b(
-            'conda', 'env', 'create', '-p', env_dir, '--file',
+            conda_exe, 'env', 'create', '-p', env_dir, '--file',
             'environment.yml', cwd=prefix.prefix_dir,
         )
         if additional_dependencies:
             cmd_output_b(
-                'conda', 'install', '-p', env_dir, *additional_dependencies,
+                conda_exe, 'install', '-p', env_dir, *additional_dependencies,
                 cwd=prefix.prefix_dir,
             )
 
diff --git a/tests/languages/conda_test.py b/tests/languages/conda_test.py
@@ -0,0 +1,38 @@
+import pytest
+
+from pre_commit import envcontext
+from pre_commit.languages.conda import _conda_exe
+
+
+@pytest.mark.parametrize(
+    ('ctx', 'expected'),
+    (
+        pytest.param(
+            (
+                ('PRE_COMMIT_USE_MICROMAMBA', envcontext.UNSET),
+                ('PRE_COMMIT_USE_MAMBA', envcontext.UNSET),
+            ),
+            'conda',
+            id='default',
+        ),
+        pytest.param(
+            (
+                ('PRE_COMMIT_USE_MICROMAMBA', '1'),
+                ('PRE_COMMIT_USE_MAMBA', ''),
+            ),
+            'micromamba',
+            id='default',
+        ),
+        pytest.param(
+            (
+                ('PRE_COMMIT_USE_MICROMAMBA', ''),
+                ('PRE_COMMIT_USE_MAMBA', '1'),
+            ),
+            'mamba',
+            id='default',
+        ),
+    ),
+)
+def test_conda_exe(ctx, expected):
+    with envcontext.envcontext(ctx):
+        assert _conda_exe() == expected
