Adding ASDF support to interpreter-search-paths

Adding [ASDF](https://asdf-vm.com) support to the python interpreter-search-paths option, through an additional special strings (<ASDF>, <ASDF_LOCAL>). ASDF actually 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]

src/python/pants/python/python_setup.py

@@ -6,8 +6,10 @@
 import logging
 import multiprocessing
 import os
+import re
+from collections import OrderedDict
 from enum import Enum
-from pathlib import Path
+from pathlib import Path, PurePath
 from typing import Iterable, List, Optional, Tuple, cast
 
 from pex.variables import Variables
@@ -112,14 +114,18 @@ def register_options(cls, register):
             "--interpreter-search-paths",
             advanced=True,
             type=list,
-            default=["<PYENV>", "<PATH>"],
+            default=["<ASDF>", "<PYENV>", "<PATH>"],
             metavar="<binary-paths>",
             help=(
                 "A list of paths to search for Python interpreters that match your project's "
                 "interpreter constraints. You can specify absolute paths to interpreter binaries "
                 "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"
@@ -209,6 +215,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),
         }
@@ -253,6 +261,102 @@ 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
+            # Just make sure that the entry is a subdirectory of the installs directory
+            if child.is_dir():
+                # And that the subdirectory has a bin directory
+                bin_dir = child / "bin"
+                if bin_dir.exists():
+                    asdf_installed_paths.append(f"{bin_dir}")
+
+        # Ignore ASDF if there are no installed versions
+        if len(asdf_installed_paths) == 0:
+            return []
+
+        asdf_paths: List[str] = []
+        asdf_versions: OrderedDict[str, bool] = 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, True) 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:
+            # This should almost ALWAYS be true for an ASDF install
+            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
+            # Contains multiple lines, one or more per tool
+            # Last line for each tool wins
+            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:
+                asdf_versions.update([(v, True) for v in re.split(r"\s+", last_line)[1:]])
+
+        for version in asdf_versions.keys():
+            install_dir = asdf_installs_dir / version / "bin"
+            if install_dir.exists():
+                asdf_paths.append(f"{install_dir}")
+            else:
+                logger.warn(
+                    f"Trying to use ASDF version {version} but {install_dir} does not exist"
+                )
+
+        # 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.
@@ -292,6 +396,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")

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, Tuple, TypeVar, Union
 
 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,57 @@ 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]) -> Tuple[_T, ...]:
+    materialized = []
+    for index in indices:
+        materialized.append(sequence[index])
+    return materialized
+
+
+@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 = materialize_indices(fake_versions, fake_home_versions)
+        fake_tool_versions.write_text(
+            f'nodejs lts\njava 8\npython {" ".join(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.append(fake_home_dir)
+        fake_dirs.append(fake_asdf_dir)
+        fake_dirs.append(fake_asdf_plugin_dir)
+        fake_dirs.append(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 +147,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: List[int] = [0, 1, 2]
+    asdf_local_versions: List[int] = [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 +238,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 +265,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",