refactor the non_host part for not injecting to custom env (#346)

fix: #343

We can get the sitepackages directory from the given interpreter
directly instead of running our code again with it.

Author: xiacunshun

src/pipdeptree/__main__.py

@@ -8,19 +8,17 @@
 from pipdeptree._cli import get_options
 from pipdeptree._discovery import get_installed_distributions
 from pipdeptree._models import PackageDAG
-from pipdeptree._non_host import handle_non_host_target
 from pipdeptree._render import render
 from pipdeptree._validate import validate
 
 
 def main(args: Sequence[str] | None = None) -> None | int:
     """CLI - The main function called as entry point."""
     options = get_options(args)
-    result = handle_non_host_target(options)
-    if result is not None:
-        return result
 
-    pkgs = get_installed_distributions(local_only=options.local_only, user_only=options.user_only)
+    pkgs = get_installed_distributions(
+        interpreter=options.python, local_only=options.local_only, user_only=options.user_only
+    )
     tree = PackageDAG.from_pkgs(pkgs)
     is_text_output = not any([options.json, options.json_tree, options.output_format])
 

src/pipdeptree/_discovery.py

@@ -1,26 +1,44 @@
 from __future__ import annotations
 
+import ast
 import site
+import subprocess  # noqa: S404
 import sys
 from importlib.metadata import Distribution, distributions
+from pathlib import Path
 from typing import Iterable, Tuple
 
 from packaging.utils import canonicalize_name
 
 
 def get_installed_distributions(
+    interpreter: str = str(sys.executable),
     local_only: bool = False,  # noqa: FBT001, FBT002
     user_only: bool = False,  # noqa: FBT001, FBT002
 ) -> list[Distribution]:
     # See https://docs.python.org/3/library/venv.html#how-venvs-work for more details.
     in_venv = sys.prefix != sys.base_prefix
     original_dists: Iterable[Distribution] = []
+    py_path = Path(interpreter).absolute()
+    using_custom_interpreter = py_path != Path(sys.executable).absolute()
 
-    if local_only and in_venv:
-        venv_site_packages = site.getsitepackages([sys.prefix])
-        original_dists = distributions(path=venv_site_packages)
-    elif user_only:
+    if user_only:
         original_dists = distributions(path=[site.getusersitepackages()])
+    elif using_custom_interpreter:
+        # We query the interpreter directly to get its `sys.path` list to be used by `distributions()`.
+        # If --python and --local-only are given, we ensure that we are only using paths associated to the interpreter's
+        # environment.
+        if local_only:
+            cmd = "import sys; print([p for p in sys.path if p.startswith(sys.prefix)])"
+        else:
+            cmd = "import sys; print(sys.path)"
+
+        args = [str(py_path), "-c", cmd]
+        result = subprocess.run(args, stdout=subprocess.PIPE, check=False, text=True)  # noqa: S603
+        original_dists = distributions(path=ast.literal_eval(result.stdout))
+    elif local_only and in_venv:
+        venv_site_packages = [p for p in sys.path if p.startswith(sys.prefix)]
+        original_dists = distributions(path=venv_site_packages)
     else:
         original_dists = distributions()
 

src/pipdeptree/_non_host.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import sys
 from importlib.metadata import PackageNotFoundError
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -26,11 +27,13 @@ def test_guess_version_setuptools(mocker: MockerFixture) -> None:
     assert result == "?"
 
 
-def test_package_as_frozen_repr(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+def test_package_as_frozen_repr(tmp_path: Path, mocker: MockerFixture) -> None:
     file_path = tmp_path / "foo.egg-link"
     with Path(file_path).open("w") as f:
         f.write("/A/B/foo")
-    monkeypatch.syspath_prepend(str(tmp_path))
+    mock_path = sys.path.copy()
+    mock_path.append(str(tmp_path))
+    mocker.patch("pipdeptree._discovery.sys.path", mock_path)
     json_text = '{"dir_info": {"editable": true}}'
     foo = Mock(metadata={"Name": "foo"}, version="20.4.1")
     foo.read_text = Mock(return_value=json_text)

tests/_models/test_package.py

@@ -16,7 +16,7 @@
     from pytest_mock import MockerFixture
 
 
-def test_local_only(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str]) -> None:
+def test_local_only(tmp_path: Path, mocker: MockerFixture, capfd: pytest.CaptureFixture[str]) -> None:
     venv_path = str(tmp_path / "venv")
     result = virtualenv.cli_run([venv_path, "--activators", ""])
     venv_site_packages = site.getsitepackages([venv_path])
@@ -27,8 +27,11 @@ def test_local_only(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capfd: pyte
         f.write("Metadata-Version: 2.3\n" "Name: foo\n" "Version: 1.2.5\n")
 
     cmd = [str(result.creator.exe.parent / "python3"), "--local-only"]
-    monkeypatch.setattr(sys, "prefix", venv_path)
-    monkeypatch.setattr(sys, "argv", cmd)
+    mocker.patch("pipdeptree._discovery.sys.prefix", venv_path)
+    sys_path = sys.path.copy()
+    mock_path = sys_path + venv_site_packages
+    mocker.patch("pipdeptree._discovery.sys.path", mock_path)
+    mocker.patch("pipdeptree._discovery.sys.argv", cmd)
     main()
     out, _ = capfd.readouterr()
     found = {i.split("==")[0] for i in out.splitlines()}
@@ -39,16 +42,16 @@ def test_local_only(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capfd: pyte
     assert found == expected
 
 
-def test_user_only(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str]) -> None:
+def test_user_only(tmp_path: Path, mocker: MockerFixture, capfd: pytest.CaptureFixture[str]) -> None:
     fake_dist = Path(tmp_path) / "foo-1.2.5.dist-info"
     fake_dist.mkdir()
     fake_metadata = Path(fake_dist) / "METADATA"
     with Path(fake_metadata).open("w") as f:
         f.write("Metadata-Version: 2.3\n" "Name: foo\n" "Version: 1.2.5\n")
 
-    monkeypatch.setattr(site, "getusersitepackages", Mock(return_value=str(tmp_path)))
     cmd = [sys.executable, "--user-only"]
-    monkeypatch.setattr(sys, "argv", cmd)
+    mocker.patch("pipdeptree._discovery.site.getusersitepackages", Mock(return_value=str(tmp_path)))
+    mocker.patch("pipdeptree._discovery.sys.argv", cmd)
     main()
     out, _ = capfd.readouterr()
     found = {i.split("==")[0] for i in out.splitlines()}

tests/test_discovery.py

@@ -12,10 +12,13 @@
 if TYPE_CHECKING:
     from pathlib import Path
 
+    from pytest_mock import MockerFixture
+
 
 @pytest.mark.parametrize("args_joined", [True, False])
 def test_custom_interpreter(
     tmp_path: Path,
+    mocker: MockerFixture,
     monkeypatch: pytest.MonkeyPatch,
     capfd: pytest.CaptureFixture[str],
     args_joined: bool,
@@ -25,7 +28,7 @@ def test_custom_interpreter(
     monkeypatch.chdir(tmp_path)
     py = str(result.creator.exe.relative_to(tmp_path))
     cmd += [f"--python={result.creator.exe}"] if args_joined else ["--python", py]
-    monkeypatch.setattr(sys, "argv", cmd)
+    mocker.patch("pipdeptree._discovery.sys.argv", cmd)
     main()
     out, _ = capfd.readouterr()
     found = {i.split("==")[0] for i in out.splitlines()}
@@ -40,10 +43,29 @@ def test_custom_interpreter(
         expected -= {"setuptools", "wheel"}
     assert found == expected, out
 
-    monkeypatch.setattr(sys, "argv", [*cmd, "--graph-output", "something"])
-    with pytest.raises(SystemExit) as context:
-        main()
-    out, err = capfd.readouterr()
-    assert context.value.code == 1
-    assert not out
-    assert err == "graphviz functionality is not supported when querying non-host python\n"
+
+def test_custom_interpreter_with_local_only(
+    tmp_path: Path,
+    mocker: MockerFixture,
+    capfd: pytest.CaptureFixture[str],
+) -> None:
+    venv_path = str(tmp_path / "venv")
+
+    result = virtualenv.cli_run([venv_path, "--system-site-packages", "--activators", ""])
+
+    cmd = ["", f"--python={result.creator.exe}", "--local-only"]
+    mocker.patch("pipdeptree._discovery.sys.prefix", venv_path)
+    mocker.patch("pipdeptree._discovery.sys.argv", cmd)
+    main()
+    out, _ = capfd.readouterr()
+    found = {i.split("==")[0] for i in out.splitlines()}
+    implementation = python_implementation()
+    if implementation == "CPython":
+        expected = {"pip", "setuptools", "wheel"}
+    elif implementation == "PyPy":  # pragma: no cover
+        expected = {"cffi", "greenlet", "pip", "readline", "setuptools", "wheel"}  # pragma: no cover
+    else:
+        raise ValueError(implementation)  # pragma: no cover
+    if sys.version_info >= (3, 12):
+        expected -= {"setuptools", "wheel"}  # pragma: no cover
+    assert found == expected, out