diff --git a/scripts/pyi_generator.py b/scripts/pyi_generator.py index a1a24859bd..143b0ce3cb 100644 --- a/scripts/pyi_generator.py +++ b/scripts/pyi_generator.py @@ -11,6 +11,7 @@ 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 @@ def _scan_file(self, module_path: Path): ) 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 @@ def scan_all(self, targets): 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():