install: Support symlinking .pyc files

Signed-off-by: Michał Górny <mgorny@gentoo.org>

Author: mgorny

gpep517/__main__.py

@@ -157,6 +157,11 @@ def add_install_args(parser):
     group.add_argument("--overwrite",
                        action="store_true",
                        help="Permit overwriting files in destdir")
+    group.add_argument("--symlink-pyc",
+                       action="store_true",
+                       help="Symlink .pyc files between optimization levels "
+                       "if their contents match")
+
     group.add_argument("--symlink-to",
                        type=PurePath,
                        help="Install symlinks to another directory rather "

gpep517/install.py

@@ -3,6 +3,7 @@
 
 import filecmp
 import functools
+import importlib.util
 import os
 import os.path
 import typing
@@ -24,12 +25,14 @@ def __init__(
             self,
             *args,
             overwrite: bool,
+            symlink_pyc: bool,
             symlink_to: typing.Optional[PurePath],
             **kwargs,
         ) -> None:
             super().__init__(*args, **kwargs)
 
             self.overwrite = overwrite
+            self.symlink_pyc = symlink_pyc
             self.symlink_to = symlink_to
 
             if symlink_to is not None:
@@ -101,6 +104,32 @@ def write_to_fs(
 
             return ret
 
+        def _compile_bytecode(self,
+                              scheme: Scheme,
+                              record: RecordEntry,
+                              ) -> None:
+            super()._compile_bytecode(scheme, record)
+
+            if not self.symlink_pyc:
+                return
+            if scheme not in ("purelib", "platlib"):
+                return
+            if not record.path.endswith(".py"):
+                return
+
+            relative_path = PurePath(self.scheme_dict[scheme]) / record.path
+            py = Path(self.destdir).joinpath(
+                relative_path.relative_to(relative_path.anchor))
+            pycs = [
+                Path(importlib.util.cache_from_source(
+                    py, optimization=opt if opt >= 1 else ""))
+                for opt in self.bytecode_optimization_levels
+            ]
+            for pyc1, pyc2 in zip(pycs, pycs[1:]):
+                if filecmp.cmp(pyc1, pyc2):
+                    pyc2.unlink()
+                    pyc2.symlink_to(pyc1.name)
+
     with WheelFile.open(wheel) as source:
         dest = DeduplicatingDestination(
             install_scheme_dict(args.prefix or DEFAULT_PREFIX,
@@ -110,6 +139,7 @@ def write_to_fs(
             bytecode_optimization_levels=args.optimize,
             destdir=str(args.destdir),
             overwrite=args.overwrite,
+            symlink_pyc=args.symlink_pyc,
             symlink_to=args.symlink_to,
         )
         logger.info(f"Installing {wheel} into {args.destdir}")

test/test_install.py

@@ -43,17 +43,22 @@ def all_files(top_path):
 @pytest.mark.parametrize("optimize", [None, "0", "1,2", "all"])
 @pytest.mark.parametrize("prefix", ["/usr", "/eprefix/usr"])
 @pytest.mark.parametrize("overwrite", [False, True])
+@pytest.mark.parametrize("symlink_pyc", [False, True])
 def test_install_wheel(tmp_path,
                        optimize: typing.Optional[str],
                        prefix: str,
-                       overwrite: bool):
+                       overwrite: bool,
+                       symlink_pyc: bool,
+                       ):
     args = (["", "install-wheel",
              "--destdir", str(tmp_path),
              "--interpreter", "/usr/bin/pythontest",
              "test/test-pkg/dist/test-1-py3-none-any.whl"] +
             (["--prefix", prefix] if prefix != "/usr" else []) +
             (["--optimize", optimize]
-             if optimize is not None else []))
+             if optimize is not None else []) +
+            (["--symlink-pyc"] if symlink_pyc else [])
+            )
     assert 0 == main(args)
 
     expected_overwrite = (contextlib.nullcontext() if overwrite
@@ -99,7 +104,10 @@ def test_install_wheel(tmp_path,
         pyc = importlib.util.cache_from_source(
             init_mod, optimization=opt if opt != 0 else "")
         expected[pathlib.Path(pyc)] = (
-            nonexec, False, pathlib.Path(f"/{init_mod}"))
+            nonexec,
+            # our only symlinking opportunity is 1 -> 0
+            symlink_pyc and opt == 1 and 0 in opt_levels,
+            pathlib.Path(f"/{init_mod}"))
 
     assert expected == dict(all_files(tmp_path))
 
@@ -129,14 +137,18 @@ def test_install_self(tmp_path):
                                           "remove-dir",
                                           "modify-file",
                                           ])
+@pytest.mark.parametrize("symlink_pyc", [False, True])
 def test_install_symlink_to(tmp_path,
                             optimize: typing.Optional[str],
                             modification: str,
+                            symlink_pyc: bool,
                             ) -> None:
     args = (["", "install-wheel",
              "--destdir", str(tmp_path),
              "test/symlink-pkg/dist/foo-0-py3-none-any.whl"] +
-            (["--optimize", optimize] if optimize is not None else []))
+            (["--optimize", optimize] if optimize is not None else []) +
+            (["--symlink-pyc"] if symlink_pyc else [])
+            )
     assert 0 == main(args + ["--prefix", "/first"])
 
     sitedir = pathlib.PurePath(
@@ -196,7 +208,10 @@ def test_install_symlink_to(tmp_path,
                 pyc = importlib.util.cache_from_source(
                     path, optimization=opt if opt != 0 else "")
                 expected[pathlib.Path(pyc)] = (
-                    False, False, pathlib.Path("/", path))
+                    False,
+                    symlink_pyc and
+                    (opt == 1 or (opt != 0 and path.name == "__init__.py")),
+                    pathlib.Path("/", path))
 
     if modification == "remove-files":
         del expected[pathlib.Path(f"first/{sitedir}/foo/a.py")]