debug

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -1,7 +1,9 @@
 # ruff: noqa: SLF001
 from __future__ import annotations
 
+import concurrent.futures
 import hashlib
+import multiprocessing
 import os
 import pickle
 import re
@@ -23,6 +25,8 @@
 from codeflash.models.models import CodePosition, FunctionCalledInTest, TestsInFile, TestType
 
 if TYPE_CHECKING:
+    from multiprocessing.synchronize import Lock
+
     from codeflash.verification.verification_utils import TestConfig
 
 
@@ -284,23 +288,33 @@ def discover_parameters_unittest(function_name: str) -> tuple[bool, str, str | N
 
 
 def _process_single_test_file(
-    test_file: Path, functions: list[TestsInFile], project_root_path: Path, test_framework: str
+    test_file: Path,
+    functions: list[TestsInFile],
+    project_root_path: Path,
+    test_framework: str,
+    jedi_lock: Optional[Lock] = None,
 ) -> tuple[Path, set]:
     import jedi
 
     local_function_to_test_map = set()
 
     try:
-        jedi_project = jedi.Project(path=project_root_path)
-        script = jedi.Script(path=test_file, project=jedi_project)
-        test_functions = set()
-
-        all_names = script.get_names(all_scopes=True, references=True)
-        all_defs = script.get_names(all_scopes=True, definitions=True)
-        all_names_top = script.get_names(all_scopes=True)
-
-        top_level_functions = {name.name: name for name in all_names_top if name.type == "function"}
-        top_level_classes = {name.name: name for name in all_names_top if name.type == "class"}
+        if jedi_lock is not None:
+            jedi_lock.acquire()
+        try:
+            jedi_project = jedi.Project(path=project_root_path)
+            script = jedi.Script(path=test_file, project=jedi_project)
+            test_functions = set()
+
+            all_names = script.get_names(all_scopes=True, references=True)
+            all_defs = script.get_names(all_scopes=True, definitions=True)
+            all_names_top = script.get_names(all_scopes=True)
+
+            top_level_functions = {name.name: name for name in all_names_top if name.type == "function"}
+            top_level_classes = {name.name: name for name in all_names_top if name.type == "class"}
+        finally:
+            if jedi_lock is not None:
+                jedi_lock.release()
     except Exception as e:
         logger.debug(f"Failed to get jedi script for {test_file}: {e}")
         return test_file, local_function_to_test_map
@@ -379,7 +393,13 @@ def _process_single_test_file(
             continue
 
         try:
-            definition = name.goto(follow_imports=True, follow_builtin_imports=False)
+            if jedi_lock is not None:
+                jedi_lock.acquire()
+            try:
+                definition = name.goto(follow_imports=True, follow_builtin_imports=False)
+            finally:
+                if jedi_lock is not None:
+                    jedi_lock.release()
         except Exception as e:
             logger.debug(str(e))
             continue
@@ -433,17 +453,25 @@ def process_test_files(
     test_framework = cfg.test_framework
 
     function_to_test_map = defaultdict(set)
-
-    with test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
-        progress,
-        task_id,
+    jedi_lock = multiprocessing.Manager().Lock()
+
+    with (
+        test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
+            progress,
+            task_id,
+        ),
+        concurrent.futures.ProcessPoolExecutor() as executor,
     ):
-        for test_file, functions in file_to_test_map.items():
-            _, local_results = _process_single_test_file(test_file, functions, project_root_path, test_framework)
-
+        futures = {
+            executor.submit(
+                _process_single_test_file, test_file, functions, project_root_path, test_framework, jedi_lock
+            ): (test_file, functions)
+            for test_file, functions in file_to_test_map.items()
+        }
+        for future in concurrent.futures.as_completed(futures):
+            _, local_results = future.result()
             for qualified_name, function_called_in_test in local_results:
                 function_to_test_map[qualified_name].add(function_called_in_test)
-
             progress.advance(task_id)
 
     function_to_tests_dict = {function: list(tests) for function, tests in function_to_test_map.items()}