Adding ASDF support to interpreter-search-paths (#12028)

Adding ASDF (https://asdf-vm.com) support to the python interpreter-search-paths option, through an additional special strings (`<ASDF>`, `<ASDF_LOCAL>`). ASDF wraps pyenv (and python-build) internally, but relocates and adapts its layout and configuration so it conforms to the standards it employs. The implementation avoids calling ASDF directly as it requires heavy use of an agumented shell (source $ASDF_DIR/asdf.sh) and a supported interpreter. Instead, it minimally re-implements the configuration and layout algorithm used by the tool to find the direct path to python interpreters, based on its standard mode of operation. This is to say, the implementation is ASDF "compatible", but may need adaptation as the tool and python plugin change over time (although the basic principles used are very unlikely to change in any way that would affect functionality). ASDF does provide support for many other languages/interpreters; additional tools could be added in the future (how they should be integrated, and how that would affect this plugin is TBD). [ci skip-rust] [ci skip-build-wheels]
pantsbuild · Jul 14, 2021 · 7c544fd · 7c544fd
1 parent 990be5c
commit 7c544fd
Show file tree

Hide file tree

Showing 4 changed files with 290 additions and 7 deletions.
diff --git a/.gitignore b/.gitignore
@@ -55,3 +55,5 @@ GSYMS
 GTAGS
 .mypy_cache/
 /.pants
+/.venv
+.tool-versions
diff --git a/src/python/pants/backend/python/util_rules/pex_environment.py b/src/python/pants/backend/python/util_rules/pex_environment.py
@@ -132,7 +132,7 @@ async def find_pex_python(
     global_options: GlobalOptions,
 ) -> PexEnvironment:
     pex_relevant_environment = await Get(
-        Environment, EnvironmentRequest(["PATH", "HOME", "PYENV_ROOT"])
+        Environment, EnvironmentRequest(["PATH", "HOME", "PYENV_ROOT", "ASDF_DIR", "ASDF_DATA_DIR"])
     )
     # PEX files are compatible with bootstrapping via Python 2.7 or Python 3.5+. The bootstrap
     # code will then re-exec itself if the underlying PEX user code needs a more specific python

diff --git a/src/python/pants/python/python_setup.py b/src/python/pants/python/python_setup.py
@@ -5,7 +5,9 @@
 
 import logging
 import os
-from pathlib import Path
+import re
+from collections import OrderedDict
+from pathlib import Path, PurePath
 from typing import Iterable, List, Optional, Tuple, cast
 
 from pex.variables import Variables
@@ -98,6 +100,10 @@ def register_options(cls, register):
                 "and/or to directories containing interpreter binaries. The order of entries does "
                 "not matter. The following special strings are supported:\n\n"
                 '* "<PATH>", the contents of the PATH env var\n'
+                '* "<ASDF>", all Python versions currently configured by ASDF '
+                "(asdf shell, ${HOME}/.tool-versions), with a fallback to all installed versions\n"
+                '* "<ASDF_LOCAL>", the ASDF interpreter with the version in '
+                "BUILD_ROOT/.tool-versions\n"
                 '* "<PYENV>", all Python versions under $(pyenv root)/versions\n'
                 '* "<PYENV_LOCAL>", the Pyenv interpreter with the version in '
                 "BUILD_ROOT/.python-version\n"
@@ -207,6 +213,8 @@ def expand_interpreter_search_paths(cls, interpreter_search_paths, env: Environm
         special_strings = {
             "<PEXRC>": cls.get_pex_python_paths,
             "<PATH>": lambda: cls.get_environment_paths(env),
+            "<ASDF>": lambda: cls.get_asdf_paths(env),
+            "<ASDF_LOCAL>": lambda: cls.get_asdf_paths(env, asdf_local=True),
             "<PYENV>": lambda: cls.get_pyenv_paths(env),
             "<PYENV_LOCAL>": lambda: cls.get_pyenv_paths(env, pyenv_local=True),
         }
@@ -251,6 +259,134 @@ def get_pex_python_paths():
         else:
             return []
 
+    @staticmethod
+    def get_asdf_paths(env: Environment, *, asdf_local: bool = False) -> List[str]:
+        """Returns a list of paths to Python interpreters managed by ASDF.
+
+        :param env: The environment to use to look up ASDF.
+        :param bool asdf_local: If True, only use the interpreter specified by
+                                '.tool-versions' file under `build_root`.
+        """
+        asdf_dir = get_asdf_dir(env)
+        if not asdf_dir:
+            return []
+
+        asdf_dir = Path(asdf_dir)
+
+        # Ignore ASDF if the python plugin isn't installed.
+        asdf_python_plugin = asdf_dir / "plugins" / "python"
+        if not asdf_python_plugin.exists():
+            return []
+
+        # Ignore ASDF if no python versions have ever been installed (the installs folder is
+        # missing).
+        asdf_installs_dir = asdf_dir / "installs" / "python"
+        if not asdf_installs_dir.exists():
+            return []
+
+        # Find all installed versions.
+        asdf_installed_paths: List[str] = []
+        for child in asdf_installs_dir.iterdir():
+            # Aliases, and non-cpython installs may have odd names.
+            # Make sure that the entry is a subdirectory of the installs directory.
+            if child.is_dir():
+                # Make sure that the subdirectory has a bin directory.
+                bin_dir = child / "bin"
+                if bin_dir.exists():
+                    asdf_installed_paths.append(str(bin_dir))
+
+        # Ignore ASDF if there are no installed versions.
+        if not asdf_installed_paths:
+            return []
+
+        asdf_paths: List[str] = []
+        asdf_versions: OrderedDict[str, str] = OrderedDict()
+        tool_versions_file = None
+
+        # Support "shell" based ASDF configuration
+        ASDF_PYTHON_VERSION = env.get("ASDF_PYTHON_VERSION")
+        if ASDF_PYTHON_VERSION:
+            asdf_versions.update(
+                [(v, "ASDF_PYTHON_VERSION") for v in re.split(r"\s+", ASDF_PYTHON_VERSION)]
+            )
+
+        # Target the local tool-versions file.
+        if asdf_local:
+            tool_versions_file = Path(get_buildroot(), ".tool-versions")
+            if not tool_versions_file.exists():
+                logger.warning(
+                    "No `.tool-versions` file found in the build root, but <ASDF_LOCAL> was set in"
+                    " `[python-setup].interpreter_search_paths`."
+                )
+                tool_versions_file = None
+        # Target the home directory tool-versions file.
+        else:
+            home = env.get("HOME")
+            if home:
+                tool_versions_file = Path(home) / ".tool-versions"
+                if not tool_versions_file.exists():
+                    tool_versions_file = None
+
+        if tool_versions_file:
+            # Parse the tool-versions file.
+            # A tool-versions file contains multiple lines, one or more per tool.
+            # Standardize that the last line for each tool wins.
+            #
+            # The definition of a tool-versions file can be found here:
+            # https://asdf-vm.com/#/core-configuration?id=tool-versions
+            tool_versions_lines = tool_versions_file.read_text().splitlines()
+            last_line = None
+            for line in tool_versions_lines:
+                # Find the last python line.
+                if line.lower().startswith("python"):
+                    last_line = line
+            if last_line:
+                _, _, versions = last_line.partition("python")
+                for v in re.split(r"\s+", versions.strip()):
+                    if ":" in v:
+                        key, _, value = v.partition(":")
+                        if key.lower() == "path":
+                            asdf_paths.append(value)
+                        elif key.lower() == "ref":
+                            asdf_versions[value] = str(tool_versions_file)
+                        else:
+                            logger.warning(
+                                f"Unknown version format `{v}` from ASDF configured by "
+                                "`[python-setup].interpreter_search_paths`, ignoring. This "
+                                "version will not be considered when determining which Python "
+                                f"interpreters to use. Please check that `{tool_versions_file}` "
+                                "is accurate."
+                            )
+                    elif v == "system":
+                        logger.warning(
+                            "System python set by ASDF configured by "
+                            "`[python-setup].interpreter_search_paths` is unsupported, ignoring. "
+                            "This version will not be considered when determining which Python "
+                            "interpreters to use. Please remove 'system' from "
+                            f"`{tool_versions_file}` to disable this warning."
+                        )
+                    else:
+                        asdf_versions[v] = str(tool_versions_file)
+
+        for version, source in asdf_versions.items():
+            install_dir = asdf_installs_dir / version / "bin"
+            if install_dir.exists():
+                asdf_paths.append(str(install_dir))
+            else:
+                logger.warning(
+                    f"Trying to use ASDF version `{version}` configured by "
+                    f"`[python-setup].interpreter_search_paths` but `{install_dir}` does not "
+                    "exist. This version will not be considered when determining which Python "
+                    f"interpreters to use. Please check that `{source}` is accurate."
+                )
+
+        # For non-local, if no paths have been defined, fallback to every version installed
+        if not asdf_local and len(asdf_paths) == 0:
+            # This could be appended to asdf_paths, but there isn't any reason to
+            return asdf_installed_paths
+        else:
+            return asdf_paths
+
     @staticmethod
     def get_pyenv_paths(env: Environment, *, pyenv_local: bool = False) -> List[str]:
         """Returns a list of paths to Python interpreters managed by pyenv.
@@ -290,6 +426,16 @@ def get_pyenv_paths(env: Environment, *, pyenv_local: bool = False) -> List[str]
         return paths
 
 
+def get_asdf_dir(env: Environment) -> PurePath | None:
+    """See https://asdf-vm.com/#/core-configuration?id=environment-variables."""
+    asdf_dir = env.get("ASDF_DIR", env.get("ASDF_DATA_DIR"))
+    if not asdf_dir:
+        home = env.get("HOME")
+        if home:
+            return PurePath(home) / ".asdf"
+    return PurePath(asdf_dir) if asdf_dir else None
+
+
 def get_pyenv_root(env: Environment) -> str | None:
     """See https://github.com/pyenv/pyenv#environment-variables."""
     from_env = env.get("PYENV_ROOT")

diff --git a/src/python/pants/python/python_setup_test.py b/src/python/pants/python/python_setup_test.py
@@ -3,16 +3,20 @@
 
 import os
 from contextlib import contextmanager
+from pathlib import Path, PurePath
+from typing import Iterable, List, Sequence, TypeVar
 
 import pytest
 
 from pants.base.build_environment import get_pants_cachedir
 from pants.engine.environment import Environment
-from pants.python.python_setup import PythonSetup, get_pyenv_root
+from pants.python.python_setup import PythonSetup, get_asdf_dir, get_pyenv_root
 from pants.testutil.rule_runner import RuleRunner
 from pants.util.contextutil import environment_as, temporary_dir
 from pants.util.dirutil import safe_mkdir_for
 
+_T = TypeVar("_T")
+
 
 @contextmanager
 def setup_pexrc_with_pex_python_path(interpreter_paths):
@@ -49,6 +53,51 @@ def fake_pyenv_root(fake_versions, fake_local_version):
         yield pyenv_root, fake_version_dirs, fake_local_version_dirs
 
 
+def materialize_indices(sequence: Sequence[_T], indices: Iterable[int]) -> List[_T]:
+    return [sequence[i] for i in indices]
+
+
+@contextmanager
+def fake_asdf_root(
+    fake_versions: List[str], fake_home_versions: List[int], fake_local_versions: List[int]
+):
+    with temporary_dir() as home_dir, temporary_dir() as asdf_dir:
+
+        fake_dirs: List[Path] = []
+        fake_version_dirs: List[str] = []
+
+        fake_home_dir = Path(home_dir)
+        fake_tool_versions = fake_home_dir / ".tool-versions"
+        fake_home_versions_str = " ".join(materialize_indices(fake_versions, fake_home_versions))
+        fake_tool_versions.write_text(f"nodejs lts\njava 8\npython {fake_home_versions_str}\n")
+
+        fake_asdf_dir = Path(asdf_dir)
+        fake_asdf_plugin_dir = fake_asdf_dir / "plugins" / "python"
+        fake_asdf_installs_dir = fake_asdf_dir / "installs" / "python"
+
+        fake_dirs.extend(
+            [fake_home_dir, fake_asdf_dir, fake_asdf_plugin_dir, fake_asdf_installs_dir]
+        )
+
+        for version in fake_versions:
+            fake_version_path = fake_asdf_installs_dir / version / "bin"
+            fake_version_dirs.append(f"{fake_version_path}")
+            fake_dirs.append(fake_version_path)
+
+        for fake_dir in fake_dirs:
+            fake_dir.mkdir(parents=True, exist_ok=True)
+
+        yield (
+            home_dir,
+            asdf_dir,
+            fake_version_dirs,
+            # fake_home_version_dirs
+            materialize_indices(fake_version_dirs, fake_home_versions),
+            # fake_local_version_dirs
+            materialize_indices(fake_version_dirs, fake_local_versions),
+        )
+
+
 def test_get_environment_paths() -> None:
     paths = PythonSetup.get_environment_paths(Environment({"PATH": "foo/bar:baz:/qux/quux"}))
     assert ["foo/bar", "baz", "/qux/quux"] == paths
@@ -92,12 +141,87 @@ def test_get_pyenv_paths(rule_runner: RuleRunner) -> None:
     assert expected_local_paths == local_paths
 
 
+def test_get_asdf_dir() -> None:
+    home = PurePath("♡")
+    default_root = home / ".asdf"
+    explicit_root = home / "explicit"
+
+    assert explicit_root == get_asdf_dir(Environment({"ASDF_DIR": f"{explicit_root}"}))
+    assert default_root == get_asdf_dir(Environment({"HOME": f"{home}"}))
+    assert get_asdf_dir(Environment({})) is None
+
+
+def test_get_asdf_paths(rule_runner: RuleRunner) -> None:
+    # 3.9.4 is intentionally "left out" so that it's only found if the "all installs" fallback is
+    # used
+    all_python_versions = ["2.7.14", "3.5.5", "3.7.10", "3.9.4", "3.9.5"]
+    asdf_home_versions = [0, 1, 2]
+    asdf_local_versions = [2, 1, 4]
+    asdf_local_versions_str = " ".join(
+        materialize_indices(all_python_versions, asdf_local_versions)
+    )
+    rule_runner.write_files(
+        {
+            ".tool-versions": (
+                "nodejs 16.0.1\n"
+                "java current\n"
+                f"python {asdf_local_versions_str}\n"
+                "rust 1.52.0\n"
+            )
+        }
+    )
+    with fake_asdf_root(all_python_versions, asdf_home_versions, asdf_local_versions) as (
+        home_dir,
+        asdf_dir,
+        expected_asdf_paths,
+        expected_asdf_home_paths,
+        expected_asdf_local_paths,
+    ):
+        # Check the "all installed" fallback
+        all_paths = PythonSetup.get_asdf_paths(Environment({"ASDF_DIR": asdf_dir}))
+
+        home_paths = PythonSetup.get_asdf_paths(
+            Environment({"HOME": home_dir, "ASDF_DIR": asdf_dir})
+        )
+        local_paths = PythonSetup.get_asdf_paths(
+            Environment({"HOME": home_dir, "ASDF_DIR": asdf_dir}), asdf_local=True
+        )
+
+        # The order the filesystem returns the "installed" folders is arbitrary
+        assert set(expected_asdf_paths) == set(all_paths)
+
+        # These have a fixed order defined by the `.tool-versions` file
+        assert expected_asdf_home_paths == home_paths
+        assert expected_asdf_local_paths == local_paths
+
+
 def test_expand_interpreter_search_paths(rule_runner: RuleRunner) -> None:
     local_pyenv_version = "3.5.5"
-    all_pyenv_versions = ["2.7.14", local_pyenv_version]
-    rule_runner.write_files({".python-version": f"{local_pyenv_version}\n"})
+    all_python_versions = ["2.7.14", local_pyenv_version, "3.7.10", "3.9.4", "3.9.5"]
+    asdf_home_versions = [0, 1, 2]
+    asdf_local_versions = [2, 1, 4]
+    asdf_local_versions_str = " ".join(
+        materialize_indices(all_python_versions, asdf_local_versions)
+    )
+    rule_runner.write_files(
+        {
+            ".python-version": f"{local_pyenv_version}\n",
+            ".tool-versions": (
+                "nodejs 16.0.1\n"
+                "java current\n"
+                f"python {asdf_local_versions_str}\n"
+                "rust 1.52.0\n"
+            ),
+        }
+    )
     with setup_pexrc_with_pex_python_path(["/pexrc/path1:/pexrc/path2"]):
-        with fake_pyenv_root(all_pyenv_versions, local_pyenv_version) as (
+        with fake_asdf_root(all_python_versions, asdf_home_versions, asdf_local_versions) as (
+            home_dir,
+            asdf_dir,
+            expected_asdf_paths,
+            expected_asdf_home_paths,
+            expected_asdf_local_paths,
+        ), fake_pyenv_root(all_python_versions, local_pyenv_version) as (
             pyenv_root,
             expected_pyenv_paths,
             expected_pyenv_local_paths,
@@ -108,11 +232,20 @@ def test_expand_interpreter_search_paths(rule_runner: RuleRunner) -> None:
                 "/bar",
                 "<PEXRC>",
                 "/baz",
+                "<ASDF>",
+                "<ASDF_LOCAL>",
                 "<PYENV>",
                 "<PYENV_LOCAL>",
                 "/qux",
             ]
-            env = Environment({"PATH": "/env/path1:/env/path2", "PYENV_ROOT": pyenv_root})
+            env = Environment(
+                {
+                    "HOME": home_dir,
+                    "PATH": "/env/path1:/env/path2",
+                    "PYENV_ROOT": pyenv_root,
+                    "ASDF_DIR": asdf_dir,
+                }
+            )
             expanded_paths = PythonSetup.expand_interpreter_search_paths(
                 paths,
                 env,
@@ -126,6 +259,8 @@ def test_expand_interpreter_search_paths(rule_runner: RuleRunner) -> None:
         "/pexrc/path1",
         "/pexrc/path2",
         "/baz",
+        *expected_asdf_home_paths,
+        *expected_asdf_local_paths,
         *expected_pyenv_paths,
         *expected_pyenv_local_paths,
         "/qux",
-Original file line number
+Diff line change
@@ Expand Up / @@ -55,3 +55,5 @@ GSYMS @@
     GTAGS
     .mypy_cache/
     /.pants
+    /.venv
+    .tool-versions