From 1a3a7d208fd48446de3d35e5c640378dc1dfd437 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20G=C3=B3rny?= <mgorny@gentoo.org>
Date: Sun, 5 Jun 2022 12:08:44 +0200
Subject: [PATCH] main: Implement --prefix option to override install path

Implement a --prefix option that can be used to override platbase
and base when expanding the paths from install scheme.  This can be used
to install paths to a layout using another prefix than the one used
to run installer.

Gentoo use case for this is running installer inside a venv to prepare
installation images for the real system.  The nixOS maintainers have
also indicated interest in this feature.
---
 src/installer/__main__.py | 23 ++++++++++++++++++-----
 tests/test_main.py        | 30 ++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/installer/__main__.py b/src/installer/__main__.py
index 3357ec5..51014b9 100644
--- a/src/installer/__main__.py
+++ b/src/installer/__main__.py
@@ -23,6 +23,13 @@ def _get_main_parser() -> argparse.ArgumentParser:
         type=str,
         help="destination directory (prefix to prepend to each file)",
     )
+    parser.add_argument(
+        "--prefix",
+        "-p",
+        metavar="path",
+        type=str,
+        help="override prefix to install packages to",
+    )
     parser.add_argument(
         "--compile-bytecode",
         action="append",
@@ -39,12 +46,18 @@ def _get_main_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def _get_scheme_dict(distribution_name: str) -> Dict[str, str]:
+def _get_scheme_dict(
+    distribution_name: str, prefix: Optional[str] = None
+) -> Dict[str, str]:
     """Calculate the scheme dictionary for the current Python environment."""
-    scheme_dict = sysconfig.get_paths()
+    vars = {}
+    if prefix is None:
+        installed_base = sysconfig.get_config_var("base")
+        assert installed_base
+    else:
+        vars["base"] = vars["platbase"] = installed_base = prefix
 
-    installed_base = sysconfig.get_config_var("base")
-    assert installed_base
+    scheme_dict = sysconfig.get_paths(vars=vars)
 
     # calculate 'headers' path, not currently in sysconfig - see
     # https://bugs.python.org/issue44445. This is based on what distutils does.
@@ -72,7 +85,7 @@ def _main(cli_args: Sequence[str], program: Optional[str] = None) -> None:
 
     with WheelFile.open(args.wheel) as source:
         destination = SchemeDictionaryDestination(
-            scheme_dict=_get_scheme_dict(source.distribution),
+            scheme_dict=_get_scheme_dict(source.distribution, prefix=args.prefix),
             interpreter=sys.executable,
             script_kind=get_launcher_kind(),
             bytecode_optimization_levels=bytecode_levels,
diff --git a/tests/test_main.py b/tests/test_main.py
index c6f7c40..391a13d 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,3 +1,5 @@
+import os
+
 from installer.__main__ import _get_scheme_dict as get_scheme_dict
 from installer.__main__ import _main as main
 
@@ -7,6 +9,14 @@ def test_get_scheme_dict():
     assert set(d.keys()) >= {"purelib", "platlib", "headers", "scripts", "data"}
 
 
+def test_get_scheme_dict_prefix():
+    d = get_scheme_dict(distribution_name="foo", prefix="/foo")
+    for key in ("purelib", "platlib", "headers", "scripts", "data"):
+        assert d[key].startswith(
+            f"{os.sep}foo"
+        ), f"{key} does not start with /foo: {d[key]}"
+
+
 def test_main(fancy_wheel, tmp_path):
     destdir = tmp_path / "dest"
 
@@ -23,6 +33,26 @@ def test_main(fancy_wheel, tmp_path):
     }
 
 
+def test_main_prefix(fancy_wheel, tmp_path):
+    destdir = tmp_path / "dest"
+
+    main([str(fancy_wheel), "-d", str(destdir), "-p", "/foo"], "python -m installer")
+
+    installed_py_files = list(destdir.rglob("*.py"))
+
+    for f in installed_py_files:
+        assert str(f.parent).startswith(
+            str(destdir / "foo")
+        ), f"path does not respect destdir+prefix: {f}"
+    assert {f.stem for f in installed_py_files} == {"__init__", "__main__", "data"}
+
+    installed_pyc_files = destdir.rglob("*.pyc")
+    assert {f.name.split(".")[0] for f in installed_pyc_files} == {
+        "__init__",
+        "__main__",
+    }
+
+
 def test_main_no_pyc(fancy_wheel, tmp_path):
     destdir = tmp_path / "dest"