CLI to install to Python running installer (#94)

Co-authored-by: Michał Górny <[email protected]>
Co-authored-by: Filipe Laíns <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Pradyun Gedam <[email protected]>

Author: 5 people

src/installer/__main__.py

@@ -0,0 +1,84 @@
+"""Installer CLI."""
+
+import argparse
+import os.path
+import sys
+import sysconfig
+from typing import Dict, Optional, Sequence
+
+import installer
+import installer.destinations
+import installer.sources
+import installer.utils
+
+
+def main_parser() -> argparse.ArgumentParser:
+    """Construct the main parser."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("wheel", type=str, help="wheel file to install")
+    parser.add_argument(
+        "--destdir",
+        "-d",
+        metavar="path",
+        type=str,
+        help="destination directory (prefix to prepend to each file)",
+    )
+    parser.add_argument(
+        "--compile-bytecode",
+        action="append",
+        metavar="level",
+        type=int,
+        choices=[0, 1, 2],
+        help="generate bytecode for the specified optimization level(s) (default=0, 1)",
+    )
+    parser.add_argument(
+        "--no-compile-bytecode",
+        action="store_true",
+        help="don't generate bytecode for installed modules",
+    )
+    return parser
+
+
+def get_scheme_dict(distribution_name: str) -> Dict[str, str]:
+    """Calculate the scheme dictionary for the current Python environment."""
+    scheme_dict = sysconfig.get_paths()
+
+    # calculate 'headers' path, not currently in sysconfig - see
+    # https://bugs.python.org/issue44445. This is based on what distutils does.
+    # TODO: figure out original vs normalised distribution names
+    scheme_dict["headers"] = os.path.join(
+        sysconfig.get_path(
+            "include", vars={"installed_base": sysconfig.get_config_var("base")}
+        ),
+        distribution_name,
+    )
+
+    return scheme_dict
+
+
+def main(cli_args: Sequence[str], program: Optional[str] = None) -> None:
+    """Process arguments and perform the install."""
+    parser = main_parser()
+    if program:
+        parser.prog = program
+    args = parser.parse_args(cli_args)
+
+    bytecode_levels = args.compile_bytecode
+    if args.no_compile_bytecode:
+        bytecode_levels = []
+    elif not bytecode_levels:
+        bytecode_levels = [0, 1]
+
+    with installer.sources.WheelFile.open(args.wheel) as source:
+        destination = installer.destinations.SchemeDictionaryDestination(
+            get_scheme_dict(source.distribution),
+            sys.executable,
+            installer.utils.get_launcher_kind(),
+            bytecode_optimization_levels=bytecode_levels,
+            destdir=args.destdir,
+        )
+        installer.install(source, destination, {})
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main(sys.argv[1:], "python -m installer")

src/installer/destinations.py

@@ -1,8 +1,19 @@
 """Handles all file writing and post-installation processing."""
 
+import compileall
 import io
 import os
-from typing import TYPE_CHECKING, BinaryIO, Dict, Iterable, Optional, Tuple, Union
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    BinaryIO,
+    Collection,
+    Dict,
+    Iterable,
+    Optional,
+    Tuple,
+    Union,
+)
 
 from installer.records import Hash, RecordEntry
 from installer.scripts import Script
@@ -99,6 +110,8 @@ def __init__(
         interpreter: str,
         script_kind: "LauncherKind",
         hash_algorithm: str = "sha256",
+        bytecode_optimization_levels: Collection[int] = (),
+        destdir: Optional[str] = None,
     ) -> None:
         """Construct a ``SchemeDictionaryDestination`` object.
 
@@ -108,11 +121,28 @@ def __init__(
         :param hash_algorithm: the hashing algorithm to use, which is a member
             of :any:`hashlib.algorithms_available` (ideally from
             :any:`hashlib.algorithms_guaranteed`).
+        :param bytecode_optimization_levels: Compile cached bytecode for
+            installed .py files with these optimization levels. The bytecode
+            is specific to the minor version of Python (e.g. 3.10) used to
+            generate it.
+        :param destdir: A staging directory in which to write all files. This
+            is expected to be the filesystem root at runtime, so embedded paths
+            will be written as though this was the root.
         """
         self.scheme_dict = scheme_dict
         self.interpreter = interpreter
         self.script_kind = script_kind
         self.hash_algorithm = hash_algorithm
+        self.bytecode_optimization_levels = bytecode_optimization_levels
+        self.destdir = destdir
+
+    def _path_with_destdir(self, scheme: Scheme, path: str) -> str:
+        file = os.path.join(self.scheme_dict[scheme], path)
+        if self.destdir is not None:
+            file_path = Path(file)
+            rel_path = file_path.relative_to(file_path.anchor)
+            return os.path.join(self.destdir, rel_path)
+        return file
 
     def write_to_fs(
         self,
@@ -131,7 +161,7 @@ def write_to_fs(
         - Ensures that an existing file is not being overwritten.
         - Hashes the written content, to determine the entry in the ``RECORD`` file.
         """
-        target_path = os.path.join(self.scheme_dict[scheme], path)
+        target_path = self._path_with_destdir(scheme, path)
         if os.path.exists(target_path):
             message = f"File already exists: {target_path}"
             raise FileExistsError(message)
@@ -201,20 +231,34 @@ def write_script(
                 Scheme("scripts"), script_name, stream, is_executable=True
             )
 
-            path = os.path.join(self.scheme_dict[Scheme("scripts")], script_name)
+            path = self._path_with_destdir(Scheme("scripts"), script_name)
             mode = os.stat(path).st_mode
             mode |= (mode & 0o444) >> 2
             os.chmod(path, mode)
 
             return entry
 
+    def _compile_bytecode(self, scheme: Scheme, record: RecordEntry) -> None:
+        """Compile bytecode for a single .py file."""
+        if scheme not in ("purelib", "platlib"):
+            return
+
+        target_path = self._path_with_destdir(scheme, record.path)
+        dir_path_to_embed = os.path.dirname(  # Without destdir
+            os.path.join(self.scheme_dict[scheme], record.path)
+        )
+        for level in self.bytecode_optimization_levels:
+            compileall.compile_file(
+                target_path, optimize=level, quiet=1, ddir=dir_path_to_embed
+            )
+
     def finalize_installation(
         self,
         scheme: Scheme,
         record_file_path: str,
         records: Iterable[Tuple[Scheme, RecordEntry]],
     ) -> None:
-        """Finalize installation, by writing the ``RECORD`` file.
+        """Finalize installation, by writing the ``RECORD`` file & compiling bytecode.
 
         :param scheme: scheme to write the ``RECORD`` file in
         :param record_file_path: path of the ``RECORD`` file with that scheme
@@ -230,7 +274,11 @@ def prefix_for_scheme(file_scheme: str) -> Optional[str]:
             )
             return path + "/"
 
-        with construct_record_file(records, prefix_for_scheme) as record_stream:
+        record_list = list(records)
+        with construct_record_file(record_list, prefix_for_scheme) as record_stream:
             self.write_to_fs(
                 scheme, record_file_path, record_stream, is_executable=False
             )
+
+        for scheme, record in record_list:
+            self._compile_bytecode(scheme, record)

tests/conftest.py

@@ -0,0 +1,75 @@
+import textwrap
+import zipfile
+
+import pytest
+
+
+@pytest.fixture
+def fancy_wheel(tmp_path):
+    path = tmp_path / "fancy-1.0.0-py2.py3-none-any.whl"
+    files = {
+        "fancy/": b"""""",
+        "fancy/__init__.py": b"""\
+            def main():
+                print("I'm fancy.")
+        """,
+        "fancy/__main__.py": b"""\
+            if __name__ == "__main__":
+                from . import main
+                main()
+        """,
+        "fancy-1.0.0.data/data/fancy/": b"""""",
+        "fancy-1.0.0.data/data/fancy/data.py": b"""\
+            # put me in data
+        """,
+        "fancy-1.0.0.dist-info/": b"""""",
+        "fancy-1.0.0.dist-info/top_level.txt": b"""\
+            fancy
+        """,
+        "fancy-1.0.0.dist-info/entry_points.txt": b"""\
+            [console_scripts]
+            fancy = fancy:main
+
+            [gui_scripts]
+            fancy-gui = fancy:main
+        """,
+        "fancy-1.0.0.dist-info/WHEEL": b"""\
+            Wheel-Version: 1.0
+            Generator: magic (1.0.0)
+            Root-Is-Purelib: true
+            Tag: py3-none-any
+        """,
+        "fancy-1.0.0.dist-info/METADATA": b"""\
+            Metadata-Version: 2.1
+            Name: fancy
+            Version: 1.0.0
+            Summary: A fancy package
+            Author: Agendaless Consulting
+            Author-email: [email protected]
+            License: MIT
+            Keywords: fancy amazing
+            Platform: UNKNOWN
+            Classifier: Intended Audience :: Developers
+        """,
+        # The RECORD file is indirectly validated by the WheelFile, since it only
+        # provides the items that are a part of the wheel.
+        "fancy-1.0.0.dist-info/RECORD": b"""\
+            fancy/__init__.py,,
+            fancy/__main__.py,,
+            fancy-1.0.0.data/data/fancy/data.py,,
+            fancy-1.0.0.dist-info/top_level.txt,,
+            fancy-1.0.0.dist-info/entry_points.txt,,
+            fancy-1.0.0.dist-info/WHEEL,,
+            fancy-1.0.0.dist-info/METADATA,,
+            fancy-1.0.0.dist-info/RECORD,,
+        """,
+    }
+
+    with zipfile.ZipFile(path, "w") as archive:
+        for name, indented_content in files.items():
+            archive.writestr(
+                name,
+                textwrap.dedent(indented_content.decode("utf-8")).encode("utf-8"),
+            )
+
+    return path

tests/test_main.py

@@ -0,0 +1,38 @@
+from installer.__main__ import get_scheme_dict, main
+
+
+def test_get_scheme_dict():
+    d = get_scheme_dict(distribution_name="foo")
+    assert set(d.keys()) >= {"purelib", "platlib", "headers", "scripts", "data"}
+
+
+def test_main(fancy_wheel, tmp_path):
+    destdir = tmp_path / "dest"
+
+    main([str(fancy_wheel), "-d", str(destdir)], "python -m installer")
+
+    installed_py_files = destdir.rglob("*.py")
+
+    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"
+
+    main(
+        [str(fancy_wheel), "-d", str(destdir), "--no-compile-bytecode"],
+        "python -m installer",
+    )
+
+    installed_py_files = destdir.rglob("*.py")
+
+    assert {f.stem for f in installed_py_files} == {"__init__", "__main__", "data"}
+
+    installed_pyc_files = destdir.rglob("*.pyc")
+    assert set(installed_pyc_files) == set()