feat: add mamba environment

wntrblm · Jun 11, 2021 · e0d650a · e0d650a
1 parent 24fc263
commit e0d650a
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 5 deletions.
diff --git a/nox/sessions.py b/nox/sessions.py
@@ -401,7 +401,7 @@ def conda_install(
             extraopts = ("--offline",)
 
         self._run(
-            "conda",
+            venv.conda_cmd,
             "install",
             "--yes",
             *extraopts,
@@ -550,12 +550,13 @@ def _create_venv(self) -> None:
                 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
@@ -152,6 +152,10 @@ class PassthroughEnv(ProcessEnv):
     hints about the actual env.
     """
 
+    @property
+    def conda_cmd(self) -> str:
+        return "conda"
+
     @staticmethod
     def is_offline() -> bool:
         """As of now this is only used in conda_install"""
@@ -176,23 +180,27 @@ 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,
         location: str,
         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:
@@ -201,7 +209,14 @@ def _clean_location(self) -> bool:
             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:
@@ -231,7 +246,7 @@ def create(self) -> bool:
 
             return False
 
-        cmd = ["conda", "create", "--yes", "--prefix", self.location]
+        cmd = [self.conda_cmd, "create", "--yes", "--prefix", self.location]
 
         cmd.extend(self.venv_params)