python module: add an automatic byte-compilation step

For all source `*.py` files installed via either py.install_sources() or
an `install_dir: py.get_install_dir()`, produce `*.pyc` files at install
time. Controllable via a module option.

3 files changed, 130 insertions, 4 deletions

diff --git a/mesonbuild/coredata.py b/mesonbuild/coredata.py
index c779422..6ce30c9 100644
--- a/mesonbuild/coredata.py
+++ b/mesonbuild/coredata.py
@@ -1266,6 +1266,8 @@ BUILTIN_CORE_OPTIONS: 'MutableKeyedOptionDictType' = OrderedDict([
      BuiltinOption(UserBooleanOption, 'Generate pkgconfig files as relocatable', False)),
 
     # Python module
+    (OptionKey('bytecompile', module='python'),
+     BuiltinOption(UserIntegerOption, 'Whether to compile bytecode', (-1, 2, 0))),
     (OptionKey('install_env', module='python'),
      BuiltinOption(UserComboOption, 'Which python environment to install to', 'prefix', choices=['auto', 'prefix', 'system', 'venv'])),
     (OptionKey('platlibdir', module='python'),
diff --git a/mesonbuild/modules/python.py b/mesonbuild/modules/python.py
index 162f8c5..a3868a0 100644
--- a/mesonbuild/modules/python.py
+++ b/mesonbuild/modules/python.py
@@ -13,9 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-import copy
-import os
-import shutil
+import copy, json, os, shutil
 import typing as T
 
 from . import ExtensionModule, ModuleInfo
@@ -41,7 +39,7 @@ if T.TYPE_CHECKING:
     from typing_extensions import TypedDict
 
     from . import ModuleState
-    from ..build import SharedModule, Data
+    from ..build import Build, SharedModule, Data
     from ..dependencies import Dependency
     from ..interpreter import Interpreter
     from ..interpreter.kwargs import ExtractRequired
@@ -66,6 +64,12 @@ mod_kwargs -= {'name_prefix', 'name_suffix'}
 
 
 class PythonExternalProgram(BasicPythonExternalProgram):
+
+    # This is a ClassVar instead of an instance bool, because although an
+    # installation is cached, we actually copy it, modify attributes such as pure,
+    # and return a temporary one rather than the cached object.
+    run_bytecompile: T.ClassVar[T.Dict[str, bool]] = {}
+
     def sanity(self, state: T.Optional['ModuleState'] = None) -> bool:
         ret = super().sanity()
         if ret:
@@ -216,6 +220,7 @@ class PythonInstallation(ExternalProgramHolder):
     )
     def install_sources_method(self, args: T.Tuple[T.List[T.Union[str, mesonlib.File]]],
                                kwargs: 'PyInstallKw') -> 'Data':
+        self.held_object.run_bytecompile[self.version] = True
         tag = kwargs['install_tag'] or 'python-runtime'
         pure = kwargs['pure'] if kwargs['pure'] is not None else self.pure
         install_dir = self._get_install_dir_impl(pure, kwargs['subdir'])
@@ -229,6 +234,7 @@ class PythonInstallation(ExternalProgramHolder):
     @noPosargs
     @typed_kwargs('python_installation.install_dir', _PURE_KW, _SUBDIR_KW)
     def get_install_dir_method(self, args: T.List['TYPE_var'], kwargs: 'PyInstallKw') -> str:
+        self.held_object.run_bytecompile[self.version] = True
         pure = kwargs['pure'] if kwargs['pure'] is not None else self.pure
         return self._get_install_dir_impl(pure, kwargs['subdir'])
 
@@ -297,6 +303,56 @@ class PythonModule(ExtensionModule):
             'find_installation': self.find_installation,
         })
 
+    def _get_install_scripts(self) -> T.List[mesonlib.ExecutableSerialisation]:
+        backend = self.interpreter.backend
+        ret = []
+        optlevel = self.interpreter.environment.coredata.get_option(mesonlib.OptionKey('bytecompile', module='python'))
+        if optlevel == -1:
+            return ret
+        if not any(PythonExternalProgram.run_bytecompile.values()):
+            return ret
+
+        installdata = backend.create_install_data()
+        py_files = []
+
+        def should_append(f, isdir: bool = False):
+            # This uses the install_plan decorated names to see if the original source was propagated via
+            # install_sources() or get_install_dir().
+            return f.startswith(('{py_platlib}', '{py_purelib}')) and (f.endswith('.py') or isdir)
+
+        for t in installdata.targets:
+            if should_append(t.out_name):
+                py_files.append(os.path.join(installdata.prefix, t.outdir, os.path.basename(t.fname)))
+        for d in installdata.data:
+            if should_append(d.install_path_name):
+                py_files.append(os.path.join(installdata.prefix, d.install_path))
+        for d in installdata.install_subdirs:
+            if should_append(d.install_path_name, True):
+                py_files.append(os.path.join(installdata.prefix, d.install_path))
+
+        import importlib.resources
+        pycompile = os.path.join(self.interpreter.environment.get_scratch_dir(), 'pycompile.py')
+        with open(pycompile, 'wb') as f:
+            f.write(importlib.resources.read_binary('mesonbuild.scripts', 'pycompile.py'))
+
+        for i in self.installations.values():
+            if isinstance(i, PythonExternalProgram) and i.run_bytecompile[i.info['version']]:
+                i = T.cast(PythonExternalProgram, i)
+                manifest = f'python-{i.info["version"]}-installed.json'
+                manifest_json = []
+                for f in py_files:
+                    if f.startswith((os.path.join(installdata.prefix, i.platlib), os.path.join(installdata.prefix, i.purelib))):
+                        manifest_json.append(f)
+                with open(os.path.join(self.interpreter.environment.get_scratch_dir(), manifest), 'w', encoding='utf-8') as f:
+                    json.dump(manifest_json, f)
+                cmd = i.command + [pycompile, manifest, str(optlevel)]
+                script = backend.get_executable_serialisation(cmd, verbose=True)
+                ret.append(script)
+        return ret
+
+    def postconf_hook(self, b: Build) -> None:
+        b.install_scripts.extend(self._get_install_scripts())
+
     # https://www.python.org/dev/peps/pep-0397/
     @staticmethod
     def _get_win_pythonpath(name_or_path: str) -> T.Optional[str]:
@@ -421,6 +477,7 @@ class PythonModule(ExtensionModule):
         else:
             python = copy.copy(python)
             python.pure = kwargs['pure']
+            python.run_bytecompile.setdefault(python.info['version'], False)
             return python
 
         raise mesonlib.MesonBugException('Unreachable code was reached (PythonModule.find_installation).')
diff --git a/mesonbuild/scripts/pycompile.py b/mesonbuild/scripts/pycompile.py
new file mode 100644
index 0000000..da92655
--- /dev/null
+++ b/mesonbuild/scripts/pycompile.py
@@ -0,0 +1,67 @@
+# Copyright 2016 The Meson development team
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ignore all lints for this file, since it is run by python2 as well
+
+# type: ignore
+# pylint: disable=deprecated-module
+
+import json, os, subprocess, sys
+from compileall import compile_file
+
+destdir = os.environ.get('DESTDIR')
+quiet = int(os.environ.get('MESON_INSTALL_QUIET', 0))
+
+def destdir_join(d1, d2):
+    if not d1:
+        return d2
+    # c:\destdir + c:\prefix must produce c:\destdir\prefix
+    parts = os.path.splitdrive(d2)
+    return d1 + parts[1]
+
+def compileall(files):
+    for f in files:
+        if destdir is not None:
+            ddir = os.path.dirname(f)
+            fullpath = destdir_join(destdir, f)
+        else:
+            ddir = None
+            fullpath = f
+
+        if os.path.isdir(fullpath):
+            for root, _, files in os.walk(fullpath):
+                ddir = os.path.dirname(os.path.splitdrive(f)[0] + root[len(destdir):])
+                for dirf in files:
+                    if dirf.endswith('.py'):
+                        fullpath = os.path.join(root, dirf)
+                        compile_file(fullpath, ddir, force=True, quiet=quiet)
+        else:
+            compile_file(fullpath, ddir, force=True, quiet=quiet)
+
+def run(manifest):
+    data_file = os.path.join(os.path.dirname(__file__), manifest)
+    with open(data_file, 'rb') as f:
+        dat = json.load(f)
+    compileall(dat)
+
+if __name__ == '__main__':
+    manifest = sys.argv[1]
+    run(manifest)
+    if len(sys.argv) > 2:
+        optlevel = int(sys.argv[2])
+        # python2 only needs one or the other
+        if optlevel == 1 or (sys.version_info >= (3,) and optlevel > 0):
+            subprocess.check_call([sys.executable, '-O'] + sys.argv[:2])
+        if optlevel == 2:
+            subprocess.check_call([sys.executable, '-OO'] + sys.argv[:2])