parallelize preprocessing across files via ProcessPoolExecutor

Replace serial file-read + tree-sitter parse loop with _preprocess_files()
using ProcessPoolExecutor. Module-level _preprocess_file() for pickling.

Covers all three callers: analyze_clusters, analyze_by_package, infer.
This commit is contained in:
tobjend 2026-07-04 02:58:47 +02:00
parent 710da56916
commit 5906adfc95

View file

@ -204,6 +204,32 @@ def frequency_filter(sequences, min_coverage=0.2):
return filtered return filtered
def _preprocess_file(fp):
"""Preprocess one file. Module-level for ProcessPoolExecutor."""
with open(fp) as f:
code = f.read()
sequences = []
for method_seq in preprocess_by_method(fp, code):
if method_seq:
sequences.append(method_seq)
return (fp, sequences)
def _preprocess_files(file_paths):
"""Preprocess multiple files in parallel."""
sequences = []
seq_files = []
n_workers = os.cpu_count()
with ProcessPoolExecutor(max_workers=n_workers) as ex:
futures = {ex.submit(_preprocess_file, fp): fp for fp in file_paths}
for f in as_completed(futures):
fp, method_seqs = f.result()
for seq in method_seqs:
sequences.append(seq)
seq_files.append(fp)
return sequences, seq_files
def analyze_clusters(file_paths, extension, project_root="", min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, include_kore=False): def analyze_clusters(file_paths, extension, project_root="", min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, include_kore=False):
"""Run full pipeline: preprocess → frequency filter → ensemble infer. """Run full pipeline: preprocess → frequency filter → ensemble infer.
@ -211,16 +237,7 @@ def analyze_clusters(file_paths, extension, project_root="", min_coverage=DEFAUL
list of (label, ensemble_result_dict, method_count, meta) tuples. list of (label, ensemble_result_dict, method_count, meta) tuples.
meta = {"files": set(paths), "imports": [lines], "arg_patterns": {...}, "packages": [...]}. meta = {"files": set(paths), "imports": [lines], "arg_patterns": {...}, "packages": [...]}.
""" """
sequences = [] sequences, seq_files = _preprocess_files(file_paths)
seq_files = []
for fp in file_paths:
with open(fp) as f:
code = f.read()
for method_seq in preprocess_by_method(fp, code):
if method_seq:
sequences.append(method_seq)
seq_files.append(fp)
if not sequences: if not sequences:
return [] return []
@ -259,17 +276,8 @@ def analyze_by_package(file_paths, extension, project_root="", min_coverage=DEFA
Returns: Returns:
list of (package_label, ensemble_result_dict, method_count, meta). list of (package_label, ensemble_result_dict, method_count, meta).
""" """
sequences = []
seq_files = []
t0 = time.time() t0 = time.time()
for fp in file_paths: sequences, seq_files = _preprocess_files(file_paths)
with open(fp) as f:
code = f.read()
for method_seq in preprocess_by_method(fp, code):
if method_seq:
sequences.append(method_seq)
seq_files.append(fp)
if not sequences: if not sequences:
return [] return []
_vprint(f"Preprocess: {len(sequences)} methods from {len(file_paths)} {extension} files ({time.time()-t0:.1f}s)") _vprint(f"Preprocess: {len(sequences)} methods from {len(file_paths)} {extension} files ({time.time()-t0:.1f}s)")
@ -326,14 +334,7 @@ def infer(file_paths, extension, min_coverage=DEFAULT_COVERAGE, prefer=None, kma
Returns: Returns:
Ensemble result dict from infer_ensemble. Ensemble result dict from infer_ensemble.
""" """
sequences = [] sequences, _ = _preprocess_files(file_paths)
for fp in file_paths:
with open(fp) as f:
code = f.read()
for method_seq in preprocess_by_method(fp, code):
if method_seq:
sequences.append(method_seq)
sequences = frequency_filter(sequences, min_coverage=0.2) sequences = frequency_filter(sequences, min_coverage=0.2)
symbol_seqs = [[text for _, text, _ in seq] for seq in sequences] symbol_seqs = [[text for _, text, _ in seq] for seq in sequences]