main: implement bytecode generation

Signed-off-by: Filipe Laíns <[email protected]>

Author: FFY00

src/installer/__main__.py

@@ -1,9 +1,11 @@
 """Installer CLI."""
 
 import argparse
+import compileall
 import distutils.dist
 import os
 import os.path
+import py_compile
 import sys
 import sysconfig
 
@@ -14,9 +16,10 @@
 from installer._compat.typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Dict, Optional, Sequence
+    from typing import Any, Dict, List, Optional, Sequence, Union
 
     from installer._compat.typing import Text
+    from installer.records import RecordEntry
 
 
 def main_parser():  # type: () -> argparse.ArgumentParser
@@ -31,6 +34,23 @@ def main_parser():  # type: () -> argparse.ArgumentParser
         default="/",
         help="destination directory",
     )
+    if sys.version_info >= (3,):
+        parser.add_argument(
+            "--optimize",
+            "-o",
+            nargs="*",
+            metavar="level",
+            type=int,
+            default=[0, 1],
+            help="optimization level(s) (default=0, 1)",
+        )
+    else:
+        parser.add_argument(
+            "--optimize",
+            "-o",
+            action="store_true",
+            help="enable optimization",
+        )
     return parser
 
 
@@ -53,6 +73,35 @@ def get_scheme_dict(distribution_name):  # type: (Text) -> Dict[str, str]
     return scheme_dict
 
 
+def _compile_records(record_dict, scheme_dict, compile_args):
+    # type: (Dict[str, List[RecordEntry]], Dict[str, str], Dict[str, Any]) -> None
+    for scheme, records in record_dict.items():
+        if scheme not in ("purelib", "platlib"):
+            continue
+        for record in records:
+            target_path = os.path.join(scheme_dict[scheme], record.path)
+            compileall.compile_file(target_path, **compile_args)
+
+
+def generate_bytecode(record_dict, scheme_dict, levels, stripdir=None):
+    # type: (Dict[str, List[RecordEntry]], Dict[str, str], Union[List[int], bool], Optional[str]) -> None
+    """Generate bytecode for Python files."""
+    if sys.version_info[0] == 2:
+        return _compile_records(record_dict, scheme_dict, {})
+
+    assert isinstance(levels, list)
+
+    compile_args = {}  # type: Dict[str, Any]
+    if stripdir:
+        compile_args["stripdir"] = stripdir
+    if sys.version_info >= (3, 7):
+        compile_args["invalidation_mode"] = py_compile.PycInvalidationMode.CHECKED_HASH
+
+    for level in levels:
+        compile_args["optimize"] = level
+        _compile_records(record_dict, scheme_dict, compile_args)
+
+
 def main(cli_args, program=None):
     # type: (Sequence[str], Optional[str]) -> None
     """Process arguments and perform the install."""
@@ -61,6 +110,8 @@ def main(cli_args, program=None):
         parser.prog = program
     args = parser.parse_args(cli_args)
 
+    bytecode_stripdir = None  # type: Optional[str]
+
     with installer.sources.WheelFile.open(args.wheel) as source:
         scheme_dict = get_scheme_dict(source.distribution)
 
@@ -76,13 +127,17 @@ def main(cli_args, program=None):
                 name: os.path.join(args.destdir, os.path.relpath(value, root))
                 for name, value in scheme_dict.items()
             }
+            bytecode_stripdir = root
 
         destination = installer.destinations.SchemeDictionaryDestination(
             scheme_dict,
             sys.executable,
             installer.utils.get_launcher_kind(),
         )
-        installer.install(source, destination, {})
+        records = installer.install(source, destination, {})
+
+    if args.optimize:
+        generate_bytecode(records, scheme_dict, args.optimize, bytecode_stripdir)
 
 
 def entrypoint():  # type: () -> None