laipeng/reflex
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

use multiprocessing to speed up pyi generation (#2453)

Browse Source
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
1 changed files with 16 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
scripts/pyi_generator.py
View File

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