ENH: Add py-rattler

CHANGES.rst

@@ -1,18 +1,3 @@
-0.6.5 (2024-08-12)
-------------------
-
-Bug Fixes
-^^^^^^^^^
-
-- Multiple Python versions are handled correctly (#1444)
-- JSONC fixes (#1426)
-
-Other Changes and Additions
-^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-- New documentation design
-
-
 0.6.4 (2024-08-12)
 ------------------
 

asv/plugin_manager.py

@@ -7,6 +7,7 @@
 from . import commands, plugins
 from .console import log
 
+ENV_PLUGINS = [".mamba", ".virtualenv", ".conda", ".rattler"]
 
 class PluginManager:
     """
@@ -30,7 +31,7 @@ def load_plugins(self, package):
                 self.init_plugin(mod)
                 self._plugins.append(mod)
             except ModuleNotFoundError as err:
-                if any(keyword in name for keyword in [".mamba", ".virtualenv", ".conda"]):
+                if any(keyword in name for keyword in ENV_PLUGINS):
                     continue  # Fine to not have these
                 else:
                     log.error(f"Couldn't load {name} because\n{err}")

asv/plugins/rattler.py

@@ -0,0 +1,141 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+import os
+import asyncio
+from pathlib import Path
+
+from yaml import load
+
+try:
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml import Loader
+
+from rattler import solve, install, VirtualPackage
+
+from .. import environment, util
+from ..console import log
+
+
+class Rattler(environment.Environment):
+    """
+    Manage an environment using py-rattler.
+
+    Dependencies are installed using py-rattler. The benchmarked
+    project is installed using the build command specified.
+    """
+
+    tool_name = "rattler"
+
+    def __init__(self, conf, python, requirements, tagged_env_vars):
+        """
+        Parameters
+        ----------
+        conf : Config instance
+
+        python : str
+            Version of Python.  Must be of the form "MAJOR.MINOR".
+
+        requirements : dict
+            Dictionary mapping a PyPI package name to a version
+            identifier string.
+        """
+        self._python = python
+        self._requirements = requirements
+        self._channels = conf.conda_channels
+        self._environment_file = None
+
+        if conf.conda_environment_file == "IGNORE":
+            log.debug(
+                "Skipping environment file due to conda_environment_file set to IGNORE"
+            )
+            self._environment_file = None
+        elif not conf.conda_environment_file:
+            if (Path("environment.yml")).exists():
+                log.debug("Using environment.yml")
+                self._environment_file = "environment.yml"
+        else:
+            if (Path(conf.conda_environment_file)).exists():
+                log.debug(f"Using {conf.conda_environment_file}")
+                self._environment_file = conf.conda_environment_file
+            else:
+                log.debug(
+                    f"Environment file {conf.conda_environment_file} not found, ignoring"
+                )
+
+        super(Rattler, self).__init__(conf, python, requirements, tagged_env_vars)
+        # Rattler configuration things
+        self._pkg_cache = f"{self._env_dir}/pkgs"
+
+        # TODO(haozeke): Provide channel priority, see mamba
+
+    def _setup(self):
+        asyncio.run(self._async_setup())
+
+    async def _async_setup(self):
+        log.info(f"Creating environment for {self.name}")
+
+        _args, pip_args = self._get_requirements()
+        _pkgs = ["python", "wheel", "pip"]  # baseline, overwritten by env file
+        env = dict(os.environ)
+        env.update(self.build_env_vars)
+        if self._environment_file:
+            # For named environments
+            env_file_name = self._environment_file
+            env_data = load(Path(env_file_name).open(), Loader=Loader)
+            _pkgs = [x for x in env_data.get("dependencies", []) if isinstance(x, str)]
+            self._channels += [
+                x for x in env_data.get("channels", []) if isinstance(x, str)
+            ]
+            self._channels = list(dict.fromkeys(self._channels).keys())
+            # Handle possible pip keys
+            pip_maybe = [
+                x for x in env_data.get("dependencies", []) if isinstance(x, dict)
+            ]
+            if len(pip_maybe) == 1:
+                try:
+                    pip_args += pip_maybe[0]["pip"]
+                except KeyError:
+                    raise KeyError("Only pip is supported as a secondary key")
+        _pkgs += _args
+        _pkgs = [util.replace_python_version(pkg, self._python) for pkg in _pkgs]
+        solved_records = await solve(
+            # Channels to use for solving
+            channels=self._channels,
+            # The specs to solve for
+            specs=_pkgs,
+            # Virtual packages define the specifications of the environment
+            virtual_packages=VirtualPackage.detect(),
+        )
+        await install(records=solved_records, target_prefix=self._path)
+        if pip_args:
+            for declaration in pip_args:
+                parsed_declaration = util.ParsedPipDeclaration(declaration)
+                pip_call = util.construct_pip_call(self._run_pip, parsed_declaration)
+                pip_call()
+
+    def _get_requirements(self):
+        _args = []
+        pip_args = []
+
+        for key, val in {**self._requirements, **self._base_requirements}.items():
+            if key.startswith("pip+"):
+                pip_args.append(f"{key[4:]} {val}")
+            else:
+                if val:
+                    _args.append(f"{key}={val}")
+                else:
+                    _args.append(key)
+
+        return _args, pip_args
+
+    def run_executable(self, executable, args, **kwargs):
+        return super(Rattler, self).run_executable(executable, args, **kwargs)
+
+    def run(self, args, **kwargs):
+        log.debug(f"Running '{' '.join(args)}' in {self.name}")
+        return self.run_executable("python", args, **kwargs)
+
+    def _run_pip(self, args, **kwargs):
+        # Run pip via python -m pip, so that it works on Windows when
+        # upgrading pip itself, and avoids shebang length limit on Linux
+        return self.run_executable("python", ["-mpip"] + list(args), **kwargs)

changelog.d/+cdf9af28.misc.rst

@@ -0,0 +1 @@
+New documentation design

changelog.d/1426.bugfix.rst

@@ -0,0 +1 @@
+JSONC fixes

changelog.d/1444.bugfix.rst

@@ -0,0 +1 @@
+Multiple python versions are now handled correctly

pyproject.toml

@@ -81,10 +81,13 @@ hg = [
 plugs = [
     "asv-bench-memray",
 ]
+envs = [
+    "py-rattler",
+]
 testR = [
     "rpy2; platform_system != 'Windows' and platform_python_implementation != 'PyPy'",
 ]
-all = ["asv[doc,dev,hg,plugs]"]
+all = ["asv[doc,dev,hg,envs]"]
 [build-system]
 requires = [
     "wheel",

test/test_environment_bench.py

@@ -13,6 +13,8 @@
     ENVIRONMENTS.append("conda")
 if tools.HAS_MAMBA:
     ENVIRONMENTS.append("mamba")
+if tools.HAS_RATTLER:
+    ENVIRONMENTS.append("rattler")
 if len(ENVIRONMENTS) == 0:
     pytest.skip("No environments can be constructed", allow_module_level=True)
 
@@ -128,26 +130,46 @@ def test_asv_benchmark(asv_project_factory, env):
 
 
 @pytest.mark.parametrize(
-    "config_modifier, expected_success, expected_error",
+    ("environment", "config_modifier", "expected_success", "expected_error"),
     [
         pytest.param(
+            env,
             {"conda_channels": ["conda-forge", "nodefaults"]},
             True,
             None,
-            id="with_conda_forge",
-        ),
+            id=f"with_conda_forge_{env}",
+            marks=[
+                pytest.mark.skipif(
+                    env == "mamba" and not tools.HAS_MAMBA, reason="needs mamba"
+                ),
+                pytest.mark.skipif(
+                    env == "rattler" and not tools.HAS_RATTLER, reason="needs rattler"
+                ),
+            ],
+        )
+        for env in ["mamba", "rattler"]
+    ]
+    + [
         pytest.param(
+            env,
             {"conda_channels": []},
             False,
-            "Solver could not find solution",
-            id="empty_conda_channels",
-        ),
+            ["Solver could not find solution", "Cannot solve the request"],
+            id=f"empty_conda_channels_{env}",
+            marks=[
+                pytest.mark.skipif(
+                    env == "mamba" and not tools.HAS_MAMBA, reason="needs mamba"
+                ),
+                pytest.mark.skipif(
+                    env == "rattler" and not tools.HAS_RATTLER, reason="needs rattler"
+                ),
+            ],
+        )
+        for env in ["mamba", "rattler"]
     ],
 )
-@pytest.mark.skipif(not tools.HAS_MAMBA,
-                    reason="needs mamba")
 def test_asv_mamba(
-    asv_project_factory, config_modifier, expected_success, expected_error
+    environment, asv_project_factory, config_modifier, expected_success, expected_error
 ):
     """
     Test running ASV benchmarks with various configurations,
@@ -156,7 +178,7 @@ def test_asv_mamba(
     project_dir = asv_project_factory(custom_config=config_modifier)
     try:
         subprocess.run(
-            ["asv", "run", "--quick", "--dry-run", "--environment", "mamba"],
+            ["asv", "run", "--quick", "--dry-run", "--environment", environment],
             cwd=project_dir,
             check=True,
             capture_output=True,
@@ -167,12 +189,13 @@ def test_asv_mamba(
     except subprocess.CalledProcessError as exc:
         if expected_success:
             pytest.fail(f"ASV benchmark unexpectedly failed: {exc.stderr}")
-        elif expected_error and expected_error not in exc.stderr:
+        elif expected_error and all([err not in exc.stderr for err in expected_error]):
             pytest.fail(
                 f"Expected error '{expected_error}' not found in stderr: {exc.stderr}"
             )
 
 
+# TODO(haozeke): Add similar tests for rattler
 @pytest.mark.parametrize(
     "create_condarc, set_mambarc, expected_success, expected_error",
     [
@@ -199,8 +222,7 @@ def test_asv_mamba(
         ),
     ],
 )
-@pytest.mark.skipif(not tools.HAS_MAMBA,
-                    reason="needs mamba")
+@pytest.mark.skipif(not tools.HAS_MAMBA, reason="needs mamba")
 def test_asv_mamba_condarc(
     asv_project_factory,
     create_condarc,

test/tools.py

@@ -79,6 +79,13 @@ def _check_conda():
     HAS_VIRTUALENV = False
 
 
+try:
+    import rattler  # noqa F401 checking if installed
+    HAS_RATTLER = True
+except ImportError:
+    HAS_RATTLER = False
+
+
 try:
     util.which(f'python{PYTHON_VER2}')
     HAS_PYTHON_VER2 = True