Repository: theacodes/nox (python)

Issue:
conda: support for Mamba
**How would this feature be useful?**

It would be great to support `mamba`, which is a conda replacement solver that is dramatically faster. A solve of a complex package like ROOT (which is a reason to use conda in the first place) takes several minutes with conda, and a few seconds with mamba.

**Describe the solution you'd like**

Option 1: Provide a new backend: `"mamba"`. This would simply make `conda_install(...)` use mamba instead of conda. This is the best option I could think of.

Option 2: Automatically use mamba if available over conda. This is slightly too magical, and there are _minute_ difference in the solve (that shouldn't matter).

Option 3: Add a setting to select the command used for conda. I don't think there is a good reason to generalize too much - mamba is a drop-in replacement for part of conda (other parts can just use conda, as it's always available with mamba).

This was mentioned in https://github.com/theacodes/nox/issues/346#issue-695644289 as a side note, but never attempted.

Pull request description:
feat: add mamba environment
Closes #444.

Initial prototype, no docs or tests yet. Does this look like a good way to do it?

Patch:
diff --git a/docs/config.rst b/docs/config.rst
index d64d2f8..f3bb1b9 100644
--- a/docs/config.rst
+++ b/docs/config.rst
@@ -159,7 +159,7 @@ You can also specify that the virtualenv should *always* be reused instead of re
     def tests(session):
         pass
 
-You are not limited to virtualenv, there is a selection of backends you can choose from as venv, conda or virtualenv (default):
+You are not limited to virtualenv, there is a selection of backends you can choose from as venv, conda, mamba, or virtualenv (default):
 
 .. code-block:: python
 
diff --git a/docs/tutorial.rst b/docs/tutorial.rst
index c6e162d..57ba7f4 100644
--- a/docs/tutorial.rst
+++ b/docs/tutorial.rst
@@ -375,6 +375,9 @@ incompatible versions of packages already installed with conda.
     session.install("contexter", "--no-deps")
     session.install("-e", ".", "--no-deps")
 
+``"mamba"`` is also allowed as a choice for ``venv_backend``, which will
+use/require `mamba <https://github.com/mamba-org/mamba>`_ instead of conda.
+
 
 Parametrization
 ---------------
diff --git a/docs/usage.rst b/docs/usage.rst
index 8103cb2..43e9795 100644
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -110,7 +110,7 @@ Then running ``nox --session tests`` will actually run all parametrized versions
 Changing the sessions default backend
 -------------------------------------
 
-By default nox uses ``virtualenv`` as the virtual environment backend for the sessions, but it also supports ``conda`` and ``venv`` as well as no backend (passthrough to whatever python environment nox is running on). You can change the default behaviour by using ``-db <backend>`` or ``--default-venv-backend <backend>``. Supported names are ``('none', 'virtualenv', 'conda', 'venv')``.
+By default nox uses ``virtualenv`` as the virtual environment backend for the sessions, but it also supports ``conda``, ``mamba``, and ``venv`` as well as no backend (passthrough to whatever python environment nox is running on). You can change the default behaviour by using ``-db <backend>`` or ``--default-venv-backend <backend>``. Supported names are ``('none', 'virtualenv', 'conda', 'mamba', 'venv')``.
 
 .. code-block:: console
 
diff --git a/nox/sessions.py b/nox/sessions.py
index 08c9cb4..83ea167 100644
--- a/nox/sessions.py
+++ b/nox/sessions.py
@@ -421,7 +421,7 @@ class Session:
             extraopts = ("--offline",)
 
         self._run(
-            "conda",
+            venv.conda_cmd,
             "install",
             "--yes",
             *extraopts,
@@ -592,12 +592,13 @@ class SessionRunner:
                 reuse_existing=reuse_existing,
                 venv_params=self.func.venv_params,
             )
-        elif backend == "conda":
+        elif backend in {"conda", "mamba"}:
             self.venv = CondaEnv(
                 self.envdir,
                 interpreter=self.func.python,  # type: ignore
                 reuse_existing=reuse_existing,
                 venv_params=self.func.venv_params,
+                conda_cmd=backend,
             )
         elif backend == "venv":
             self.venv = VirtualEnv(
diff --git a/nox/virtualenv.py b/nox/virtualenv.py
index 1dec7d0..501118f 100644
--- a/nox/virtualenv.py
+++ b/nox/virtualenv.py
@@ -156,6 +156,8 @@ class PassthroughEnv(ProcessEnv):
     hints about the actual env.
     """
 
+    conda_cmd = "conda"
+
     @staticmethod
     def is_offline() -> bool:
         """As of now this is only used in conda_install"""
@@ -180,10 +182,11 @@ class CondaEnv(ProcessEnv):
             If not specified, this will use the currently running Python.
         reuse_existing (Optional[bool]): Flag indicating if the conda environment
             should be reused if it already exists at ``location``.
+        conda_cmd (str): The name of the command, can be "conda" (default) or "mamba".
     """
 
     is_sandboxed = True
-    allowed_globals = ("conda",)
+    allowed_globals = ("conda", "mamba")
 
     def __init__(
         self,
@@ -191,12 +194,15 @@ class CondaEnv(ProcessEnv):
         interpreter: Optional[str] = None,
         reuse_existing: bool = False,
         venv_params: Any = None,
+        *,
+        conda_cmd: str = "conda",
     ):
         self.location_name = location
         self.location = os.path.abspath(location)
         self.interpreter = interpreter
         self.reuse_existing = reuse_existing
         self.venv_params = venv_params if venv_params else []
+        self.conda_cmd = conda_cmd
         super(CondaEnv, self).__init__()
 
     def _clean_location(self) -> bool:
@@ -205,7 +211,14 @@ class CondaEnv(ProcessEnv):
             if self.reuse_existing:
                 return False
             else:
-                cmd = ["conda", "remove", "--yes", "--prefix", self.location, "--all"]
+                cmd = [
+                    self.conda_cmd,
+                    "remove",
+                    "--yes",
+                    "--prefix",
+                    self.location,
+                    "--all",
+                ]
                 nox.command.run(cmd, silent=True, log=False)
                 # Make sure that location is clean
                 try:
@@ -233,7 +246,7 @@ class CondaEnv(ProcessEnv):
 
             return False
 
-        cmd = ["conda", "create", "--yes", "--prefix", self.location]
+        cmd = [self.conda_cmd, "create", "--yes", "--prefix", self.location]
 
         cmd.extend(self.venv_params)
 


Test patch:
diff --git a/tests/test_sessions.py b/tests/test_sessions.py
index 6bcd916..a7d1642 100644
--- a/tests/test_sessions.py
+++ b/tests/test_sessions.py
@@ -384,7 +384,8 @@ class TestSession:
         "auto_offline", [False, True], ids="auto_offline={}".format
     )
     @pytest.mark.parametrize("offline", [False, True], ids="offline={}".format)
-    def test_conda_install(self, auto_offline, offline):
+    @pytest.mark.parametrize("conda", ["conda", "mamba"], ids=str)
+    def test_conda_install(self, auto_offline, offline, conda):
         runner = nox.sessions.SessionRunner(
             name="test",
             signatures=["test"],
@@ -396,6 +397,7 @@ class TestSession:
         runner.venv.location = "/path/to/conda/env"
         runner.venv.env = {}
         runner.venv.is_offline = lambda: offline
+        runner.venv.conda_cmd = conda
 
         class SessionNoSlots(nox.sessions.Session):
             pass
@@ -406,7 +408,7 @@ class TestSession:
             args = ("--offline",) if auto_offline and offline else ()
             session.conda_install("requests", "urllib3", auto_offline=auto_offline)
             run.assert_called_once_with(
-                "conda",
+                conda,
                 "install",
                 "--yes",
                 *args,
@@ -434,6 +436,7 @@ class TestSession:
         runner.venv.location = "/path/to/conda/env"
         runner.venv.env = {}
         runner.venv.is_offline = lambda: True
+        runner.venv.conda_cmd = "conda"
 
         runner.global_config.no_install = no_install
         runner.venv._reused = reused
@@ -460,6 +463,7 @@ class TestSession:
         runner.venv.location = "/path/to/conda/env"
         runner.venv.env = {}
         runner.venv.is_offline = lambda: False
+        runner.venv.conda_cmd = "conda"
 
         class SessionNoSlots(nox.sessions.Session):
             pass