Support iterating a venv's distributions.

Work towards pex-tool#1361.

Author: jsirois

pex/interpreter.py

@@ -802,7 +802,9 @@ def _spawn_from_binary(cls, binary):
         # type: (str) -> SpawnedJob[PythonInterpreter]
         canonicalized_binary = cls.canonicalize_path(binary)
         if not os.path.exists(canonicalized_binary):
-            raise cls.InterpreterNotFound(canonicalized_binary)
+            raise cls.InterpreterNotFound(
+                "The interpreter path {} does not exist.".format(canonicalized_binary)
+            )
 
         # N.B.: The cache is written as the last step in PythonInterpreter instance initialization.
         cached_interpreter = cls._PYTHON_INTERPRETER_BY_NORMALIZED_PATH.get(canonicalized_binary)

pex/tools/commands/virtualenv.py

@@ -9,7 +9,9 @@
 import re
 import sys
 from contextlib import closing
+from textwrap import dedent
 
+from pex import third_party
 from pex.common import AtomicDirectory, is_exe, safe_mkdir
 from pex.compatibility import get_stdout_bytes_buffer
 from pex.interpreter import PythonInterpreter
@@ -19,7 +21,10 @@
 from pex.util import named_temporary_file
 
 if TYPE_CHECKING:
-    from typing import Iterator, Optional
+    import attr  # vendor:skip
+    from typing import Iterator, Optional, Union
+else:
+    from pex.third_party import attr
 
 _MIN_PIP_PYTHON_VERSION = (2, 7, 9)
 
@@ -71,7 +76,32 @@ def _is_python_script(executable):
         )
 
 
+class InvalidVirtualenvError(Exception):
+    pass
+
+
+@attr.s(frozen=True)
+class DistributionInfo(object):
+    project_name = attr.ib()  # type: str
+    version = attr.ib()  # type: str
+    sys_path_entry = attr.ib()  # type: str
+
+
 class Virtualenv(object):
+    @classmethod
+    def enclosing(cls, python):
+        # type: (Union[str, PythonInterpreter]) -> Optional[Virtualenv]
+        interpreter = (
+            python
+            if isinstance(python, PythonInterpreter)
+            else PythonInterpreter.from_binary(python)
+        )
+        if not interpreter.is_venv:
+            return None
+        return cls(
+            venv_dir=interpreter.prefix, python_exe_name=os.path.basename(interpreter.binary)
+        )
+
     @classmethod
     def create(
         cls,
@@ -163,9 +193,16 @@ def __init__(
         self._venv_dir = venv_dir
         self._custom_prompt = custom_prompt
         self._bin_dir = os.path.join(venv_dir, "bin")
-        self._interpreter = PythonInterpreter.from_binary(
-            os.path.join(self._bin_dir, python_exe_name)
-        )
+        python_exe_path = os.path.join(self._bin_dir, python_exe_name)
+        try:
+            self._interpreter = PythonInterpreter.from_binary(python_exe_path)
+        except PythonInterpreter.InterpreterNotFound as e:
+            raise InvalidVirtualenvError(
+                "The virtualenv at {venv_dir} is not valid. Failed to load an interpreter at "
+                "{python_exe_path}: {err}".format(
+                    venv_dir=self._venv_dir, python_exe_path=python_exe_path, err=e
+                )
+            )
         self._site_packages_dir = (
             os.path.join(venv_dir, "site-packages")
             if self._interpreter.identity.interpreter == "PyPy"
@@ -178,6 +215,13 @@ def __init__(
                 "site-packages",
             )
         )
+        if not os.path.isdir(self._site_packages_dir):
+            raise InvalidVirtualenvError(
+                "The virtualenv at {venv_dir} is not valid. The expected site-packages directory "
+                "at {site_packages_dir} does not exist.".format(
+                    venv_dir=venv_dir, site_packages_dir=self._site_packages_dir
+                )
+            )
         self._base_bin = frozenset(_iter_files(self._bin_dir))
 
     @property
@@ -219,6 +263,48 @@ def iter_executables(self):
             if is_exe(path):
                 yield path
 
+    def iter_distributions(self):
+        # type: () -> Iterator[DistributionInfo]
+        """"""
+        setuptools_path = tuple(third_party.expose(["setuptools"]))
+        _, stdout, _ = self.interpreter.execute(
+            args=[
+                "-c",
+                dedent(
+                    """\
+                    from __future__ import print_function
+
+                    import sys
+
+
+                    setuptools_path = {setuptools_path!r}
+                    sys.path.extend(setuptools_path)
+
+                    from pkg_resources import working_set
+
+
+                    for dist in working_set:
+                        if dist.location in setuptools_path:
+                            continue
+                        print(
+                            "{{project_name}} {{version}} {{sys_path_entry}}".format(
+                                project_name=dist.project_name,
+                                version=dist.version,
+                                sys_path_entry=dist.location,
+                            )
+                        )
+                    """.format(
+                        setuptools_path=setuptools_path
+                    )
+                ),
+            ]
+        )
+        for line in stdout.splitlines():
+            project_name, version, sys_path_entry = line.split()
+            yield DistributionInfo(
+                project_name=project_name, version=version, sys_path_entry=sys_path_entry
+            )
+
     def _rewrite_base_scripts(self, real_venv_dir):
         # type: (str) -> Iterator[str]
         scripts = [

tests/integration/tools/commands/test_virtualenv.py

@@ -0,0 +1,86 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import os.path
+import shutil
+import subprocess
+
+import pytest
+
+from pex.pep_503 import ProjectName
+from pex.testing import ALL_PY_VERSIONS, ensure_python_venv
+from pex.tools.commands.virtualenv import DistributionInfo, InvalidVirtualenvError, Virtualenv
+from pex.typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Dict
+
+
+def test_invalid(tmpdir):
+    # type: (Any) -> None
+
+    with pytest.raises(InvalidVirtualenvError):
+        Virtualenv(venv_dir=str(tmpdir))
+
+    venv_dir = os.path.join(str(tmpdir), "venv")
+    Virtualenv.create(venv_dir=venv_dir)
+    venv = Virtualenv(venv_dir=venv_dir)
+
+    shutil.rmtree(venv.site_packages_dir)
+    with pytest.raises(InvalidVirtualenvError):
+        Virtualenv(venv_dir=venv_dir)
+
+
+def test_enclosing(tmpdir):
+    # type: (Any) -> None
+
+    venv_dir = os.path.join(str(tmpdir), "venv")
+    venv = Virtualenv.create(venv_dir=venv_dir)
+
+    enclosing = Virtualenv.enclosing(venv.interpreter)
+    assert enclosing is not None
+    assert venv_dir == enclosing.venv_dir
+
+    enclosing = Virtualenv.enclosing(venv.interpreter.binary)
+    assert enclosing is not None
+    assert venv_dir == enclosing.venv_dir
+
+    assert Virtualenv.enclosing(venv.interpreter.resolve_base_interpreter()) is None
+
+
+def index_distributions(venv):
+    # type: (Virtualenv) -> Dict[ProjectName, DistributionInfo]
+    return {
+        ProjectName(dist_info.project_name): dist_info for dist_info in venv.iter_distributions()
+    }
+
+
+def test_iter_distributions_empty(tmpdir):
+    # type: (Any) -> None
+
+    empty_venv_dir = os.path.join(str(tmpdir), "empty.venv")
+    empty_venv = Virtualenv.create(venv_dir=empty_venv_dir)
+    assert {} == index_distributions(empty_venv)
+
+
+@pytest.mark.parametrize("py_version", ALL_PY_VERSIONS)
+def test_iter_distributions(tmpdir, py_version):
+    # type: (Any, str) -> None
+
+    python, pip = ensure_python_venv(py_version)
+
+    venv = Virtualenv.enclosing(python)
+    assert venv is not None
+
+    dists = index_distributions(venv)
+    pip_dist_info = dists.get(ProjectName("pip"))
+    assert pip_dist_info is not None, "Expected venv to have Pip installed."
+    assert venv.site_packages_dir == pip_dist_info.sys_path_entry
+    assert ProjectName("cowsay") not in dists
+
+    subprocess.check_call(args=[pip, "install", "cowsay==4.0"])
+    dists = index_distributions(venv)
+    cowsay_dist_info = dists.get(ProjectName("cowsay"))
+    assert cowsay_dist_info is not None, "Expected venv to have cowsay installed."
+    assert "4.0" == cowsay_dist_info.version
+    assert venv.site_packages_dir == cowsay_dist_info.sys_path_entry