efficient passes

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -503,12 +503,20 @@ def process_test_files(
                 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"}
+                # Single call to get all names with references and definitions
+                all_names = script.get_names(all_scopes=True, references=True, definitions=True)
+
+                # Filter once and create lookup dictionaries
+                top_level_functions = {}
+                top_level_classes = {}
+                all_defs = []
+
+                for name in all_names:
+                    if name.type == "function":
+                        top_level_functions[name.name] = name
+                        all_defs.append(name)
+                    elif name.type == "class":
+                        top_level_classes[name.name] = name
             except Exception as e:
                 logger.debug(f"Failed to get jedi script for {test_file}: {e}")
                 progress.advance(task_id)
@@ -574,23 +582,24 @@ def process_test_files(
                                     )
 
             test_functions_list = list(test_functions)
-            test_functions_raw = [elem.function_name for elem in test_functions_list]
-
-            test_functions_by_name = defaultdict(list)
-            for i, func_name in enumerate(test_functions_raw):
-                test_functions_by_name[func_name].append(i)
-
-            for name in all_names:
-                if name.full_name is None:
-                    continue
-                m = FUNCTION_NAME_REGEX.search(name.full_name)
-                if not m:
-                    continue
-
-                scope = m.group(1)
-                if scope not in test_functions_by_name:
-                    continue
-
+            test_functions_by_name = {}
+            for i, func in enumerate(test_functions_list):
+                if func.function_name not in test_functions_by_name:
+                    test_functions_by_name[func.function_name] = []
+                test_functions_by_name[func.function_name].append((i, func))
+
+            relevant_names = [
+                (name, m.group(1))
+                for name in all_names
+                if (
+                    name.full_name is not None
+                    and (m := FUNCTION_NAME_REGEX.search(name.full_name))
+                    and m.group(1) in test_functions_by_name
+                )
+            ]
+
+            # Process relevant names in batch
+            for name, scope in relevant_names:
                 try:
                     definition = name.goto(follow_imports=True, follow_builtin_imports=False)
                 except Exception as e:
@@ -600,36 +609,36 @@ def process_test_files(
                 if not definition or definition[0].type != "function":
                     continue
 
-                definition_path = str(definition[0].module_path)
+                definition_obj = definition[0]
+                definition_path = str(definition_obj.module_path)
+
                 if (
                     definition_path.startswith(str(project_root_path) + os.sep)
-                    and definition[0].module_name != name.module_name
-                    and definition[0].full_name is not None
+                    and definition_obj.module_name != name.module_name
+                    and definition_obj.full_name is not None
                 ):
-                    for index in test_functions_by_name[scope]:
-                        scope_test_function = test_functions_list[index].function_name
-                        scope_test_class = test_functions_list[index].test_class
-                        scope_parameters = test_functions_list[index].parameters
-                        test_type = test_functions_list[index].test_type
+                    # Pre-compute common values outside the inner loop
+                    full_name_without_module_prefix = definition_obj.full_name.replace(
+                        definition_obj.module_name + ".", "", 1
+                    )
+                    qualified_name_with_modules_from_root = f"{module_name_from_file_path(definition_obj.module_path, project_root_path)}.{full_name_without_module_prefix}"
 
-                        if scope_parameters is not None:
-                            if test_framework == "pytest":
-                                scope_test_function += "[" + scope_parameters + "]"
-                            if test_framework == "unittest":
-                                scope_test_function += "_" + scope_parameters
+                    for _index, test_func in test_functions_by_name[scope]:
+                        scope_test_function = test_func.function_name
 
-                        full_name_without_module_prefix = definition[0].full_name.replace(
-                            definition[0].module_name + ".", "", 1
-                        )
-                        qualified_name_with_modules_from_root = f"{module_name_from_file_path(definition[0].module_path, project_root_path)}.{full_name_without_module_prefix}"
+                        if test_func.parameters is not None:
+                            if test_framework == "pytest":
+                                scope_test_function += "[" + test_func.parameters + "]"
+                            elif test_framework == "unittest":
+                                scope_test_function += "_" + test_func.parameters
 
                         function_to_test_map[qualified_name_with_modules_from_root].add(
                             FunctionCalledInTest(
                                 tests_in_file=TestsInFile(
                                     test_file=test_file,
-                                    test_class=scope_test_class,
+                                    test_class=test_func.test_class,
                                     test_function=scope_test_function,
-                                    test_type=test_type,
+                                    test_type=test_func.test_type,
                                 ),
                                 position=CodePosition(line_no=name.line, col_no=name.column),
                             )