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)
@@ -573,24 +581,21 @@ 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 func in test_functions:
+                test_functions_by_name[func.function_name].append(func)
 
-            for name in all_names:
-                if name.full_name is None:
-                    continue
-                m = FUNCTION_NAME_REGEX.search(name.full_name)
-                if not m:
-                    continue
+            test_function_names_set = set(test_functions_by_name.keys())
+            relevant_names = []
 
-                scope = m.group(1)
-                if scope not in test_functions_by_name:
-                    continue
+            names_with_full_name = [name for name in all_names if name.full_name is not None]
 
+            for name in names_with_full_name:
+                match = FUNCTION_NAME_REGEX.search(name.full_name)
+                if match and match.group(1) in test_function_names_set:
+                    relevant_names.append((name, match.group(1)))
+
+            for name, scope in relevant_names:
                 try:
                     definition = name.goto(follow_imports=True, follow_builtin_imports=False)
                 except Exception as e:
@@ -600,36 +605,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)
+
+                project_root_str = str(project_root_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
+                    definition_path.startswith(project_root_str + os.sep)
+                    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
+                    module_prefix = definition_obj.module_name + "."
+                    full_name_without_module_prefix = definition_obj.full_name.replace(module_prefix, "", 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:
+                    for test_func in test_functions_by_name[scope]:
+                        if test_func.parameters is not None:
                             if test_framework == "pytest":
-                                scope_test_function += "[" + scope_parameters + "]"
-                            if test_framework == "unittest":
-                                scope_test_function += "_" + scope_parameters
-
-                        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}"
+                                scope_test_function = f"{test_func.function_name}[{test_func.parameters}]"
+                            else:  # unittest
+                                scope_test_function = f"{test_func.function_name}_{test_func.parameters}"
+                        else:
+                            scope_test_function = test_func.function_name
 
                         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),
                             )