use multiprocessing to speed up pyi generation (#2453)

This commit is contained in:
benedikt-bartscher 2024-01-25 22:46:15 +01:00 committed by GitHub
parent 8ae67f89e3
commit 4219a78e7c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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():