Explorar el Código

use multiprocessing to speed up pyi generation (#2453)

benedikt-bartscher hace 1 año
padre
commit
4219a78e7c
Se han modificado 1 ficheros con 16 adiciones y 11 borrados
  1. 16 11
      scripts/pyi_generator.py

+ 16 - 11
scripts/pyi_generator.py

@@ -11,6 +11,7 @@ import sys
 import textwrap
 import typing
 from inspect import getfullargspec
+from multiprocessing import Pool, cpu_count
 from pathlib import Path
 from types import ModuleType
 from typing import Any, Callable, Iterable, Type, get_args
@@ -637,13 +638,9 @@ class PyiGenerator:
         )
         self._write_pyi_file(module_path, ast.unparse(new_tree))
 
-    def _scan_folder(self, folder):
-        for root, _, files in os.walk(folder):
-            for file in files:
-                if file in EXCLUDED_FILES:
-                    continue
-                if file.endswith(".py"):
-                    self._scan_file(Path(root) / file)
+    def _scan_files_multiprocess(self, files):
+        with Pool(processes=cpu_count()) as pool:
+            pool.map(self._scan_file, files)
 
     def scan_all(self, targets):
         """Scan all targets for class inheriting Component and generate the .pyi files.
@@ -651,11 +648,19 @@ class PyiGenerator:
         Args:
             targets: the list of file/folders to scan.
         """
+        file_targets = []
         for target in targets:
-            if target.endswith(".py"):
-                self._scan_file(Path(target))
-            else:
-                self._scan_folder(target)
+            path = Path(target)
+            if target.endswith(".py") and path.is_file():
+                file_targets.append(path)
+            elif path.is_dir():
+                for root, _, files in os.walk(path):
+                    for file in files:
+                        if file in EXCLUDED_FILES or not file.endswith(".py"):
+                            continue
+                        file_targets.append(Path(root) / file)
+
+        self._scan_files_multiprocess(file_targets)
 
 
 def generate_init():