Use compile_dir to compile

Author: notatallshaw

src/pip/_internal/operations/install/wheel.py

@@ -76,6 +76,8 @@ def save(self) -> None:
 RecordPath = NewType("RecordPath", str)
 InstalledCSVRow = Tuple[RecordPath, str, Union[int, str]]
 
+IS_NOT_PYTHON_FILE_RE = re.compile(r'^.*(?<!\.py)$')
+
 
 def rehash(path: str, blocksize: int = 1 << 20) -> Tuple[str, str]:
     """Return (encoded_digest, length) for path using hashlib.sha256()"""
@@ -596,13 +598,34 @@ def pyc_source_file_paths() -> Generator[str, None, None]:
         # Sorting installation paths makes it easier to reproduce and debug
         # issues related to permissions on existing files.
         for installed_path in sorted(set(installed.values())):
+            if not installed_path.endswith(".py"):
+                continue
+
             full_installed_path = os.path.join(lib_dir, installed_path)
             if not os.path.isfile(full_installed_path):
                 continue
-            if not full_installed_path.endswith(".py"):
-                continue
             yield full_installed_path
 
+    def pyc_source_root_dirs() -> Generator[str, None, None]:
+        # All root source directories of the install (i.e. those that
+        # are direct children of lib_dir) that the source Python files
+        # are located within. Typically this is just one directory.
+        source_root_dirs = set()
+        for installed_path in installed.values():
+            if not installed_path.endswith(".py"):
+                continue
+
+            full_installed_path = os.path.join(lib_dir, installed_path)
+            if not os.path.isfile(full_installed_path):
+                continue
+
+            parent_full_installed_path = os.path.dirname(full_installed_path)
+            if os.path.dirname(parent_full_installed_path) == lib_dir:
+                if parent_full_installed_path in source_root_dirs:
+                    continue
+                source_root_dirs.add(parent_full_installed_path)
+                yield parent_full_installed_path
+
     def pyc_output_path(path: str) -> str:
         """Return the path the pyc file would have been written to."""
         return importlib.util.cache_from_source(path)
@@ -614,11 +637,16 @@ def pyc_output_path(path: str) -> str:
         ) as stdout:
             with warnings.catch_warnings():
                 warnings.filterwarnings("ignore")
+                # Use compile_dir on the root directories in the package,
+                # as compile_dir uses multiprocess this will be quicker
+                # than using compile_file on each file.
+                for dir in pyc_source_root_dirs():
+                    compileall.compile_dir(dir, force=True, quiet=2, workers=0, rx=IS_NOT_PYTHON_FILE_RE)
+
+                # Now check indivdual files were created
                 for path in pyc_source_file_paths():
-                    success = compileall.compile_file(path, force=True, quiet=True)
-                    if success:
-                        pyc_path = pyc_output_path(path)
-                        assert os.path.exists(pyc_path)
+                    pyc_path = pyc_output_path(path)
+                    if os.path.exists(pyc_path):
                         pyc_record_path = cast(
                             "RecordPath", pyc_path.replace(os.path.sep, "/")
                         )