fix: resolve mypy errors in discovery and support files

- Change FunctionVisitor.file_path from str to Path
- Unify dict keys to Path across discovery functions (get_all_files_and_functions,
  get_functions_within_lines, get_functions_within_git_diff, etc.)
- Remove redundant isinstance check in discover_functions
- Add assert for found_function narrowing after exit_with_message
- Fix closest_matching_file_function_name return type narrowing

Author: KRRT7

codeflash/discovery/functions_to_optimize.py

@@ -68,9 +68,9 @@ def visit_Return(self, node: cst.Return) -> None:
 class FunctionVisitor(cst.CSTVisitor):
     METADATA_DEPENDENCIES = (cst.metadata.PositionProvider, cst.metadata.ParentNodeProvider)
 
-    def __init__(self, file_path: str) -> None:
+    def __init__(self, file_path: Path) -> None:
         super().__init__()
-        self.file_path: str = file_path
+        self.file_path: Path = file_path
         self.functions: list[FunctionToOptimize] = []
 
     @staticmethod
@@ -264,7 +264,7 @@ def get_functions_to_optimize(
     assert sum([bool(optimize_all), bool(replay_test), bool(file)]) <= 1, (
         "Only one of optimize_all, replay_test, or file should be provided"
     )
-    functions: dict[str, list[FunctionToOptimize]]
+    functions: dict[Path, list[FunctionToOptimize]]
     trace_file_path: Path | None = None
     is_lsp = is_LSP_enabled()
     with warnings.catch_warnings():
@@ -281,7 +281,7 @@ def get_functions_to_optimize(
             logger.info("!lsp|Finding all functions in the file '%s'…", file)
             console.rule()
             file = Path(file) if isinstance(file, str) else file
-            functions: dict[Path, list[FunctionToOptimize]] = find_all_functions_in_file(file)
+            functions = find_all_functions_in_file(file)
             if only_get_this_function is not None:
                 split_function = only_get_this_function.split(".")
                 if len(split_function) > 2:
@@ -316,6 +316,7 @@ def get_functions_to_optimize(
                         f"Function {only_get_this_function} not found in file {file}\nor the function does not have a 'return' statement or is a property"
                     )
 
+                assert found_function is not None
                 # For JavaScript/TypeScript, verify that the function (or its parent class) is exported
                 # Non-exported functions cannot be imported by tests
                 if found_function.language in ("javascript", "typescript"):
@@ -359,7 +360,7 @@ def get_functions_to_optimize(
         return filtered_modified_functions, functions_count, trace_file_path
 
 
-def get_functions_within_git_diff(uncommitted_changes: bool) -> dict[str, list[FunctionToOptimize]]:
+def get_functions_within_git_diff(uncommitted_changes: bool) -> dict[Path, list[FunctionToOptimize]]:
     modified_lines: dict[str, list[int]] = get_git_diff(uncommitted_changes=uncommitted_changes)
     return get_functions_within_lines(modified_lines)
 
@@ -400,7 +401,7 @@ def closest_matching_file_function_name(
                 closest_match = function
                 closest_file = file_path
 
-    if closest_match is not None:
+    if closest_match is not None and closest_file is not None:
         return closest_file, closest_match
     return None
 
@@ -434,13 +435,13 @@ def levenshtein_distance(s1: str, s2: str) -> int:
     return previous[len1]
 
 
-def get_functions_inside_a_commit(commit_hash: str) -> dict[str, list[FunctionToOptimize]]:
+def get_functions_inside_a_commit(commit_hash: str) -> dict[Path, list[FunctionToOptimize]]:
     modified_lines: dict[str, list[int]] = get_git_diff(only_this_commit=commit_hash)
     return get_functions_within_lines(modified_lines)
 
 
-def get_functions_within_lines(modified_lines: dict[str, list[int]]) -> dict[str, list[FunctionToOptimize]]:
-    functions: dict[str, list[FunctionToOptimize]] = {}
+def get_functions_within_lines(modified_lines: dict[str, list[int]]) -> dict[Path, list[FunctionToOptimize]]:
+    functions: dict[Path, list[FunctionToOptimize]] = {}
     for path_str, lines_in_file in modified_lines.items():
         path = Path(path_str)
         if not path.exists():
@@ -452,9 +453,9 @@ def get_functions_within_lines(modified_lines: dict[str, list[int]]) -> dict[str
             except Exception as e:
                 logger.exception(e)
                 continue
-            function_lines = FunctionVisitor(file_path=str(path))
+            function_lines = FunctionVisitor(file_path=path)
             wrapper.visit(function_lines)
-            functions[str(path)] = [
+            functions[path] = [
                 function_to_optimize
                 for function_to_optimize in function_lines.functions
                 if (start_line := function_to_optimize.starting_line) is not None
@@ -466,7 +467,7 @@ def get_functions_within_lines(modified_lines: dict[str, list[int]]) -> dict[str
 
 def get_all_files_and_functions(
     module_root_path: Path, ignore_paths: list[Path], language: Language | None = None
-) -> dict[str, list[FunctionToOptimize]]:
+) -> dict[Path, list[FunctionToOptimize]]:
     """Get all optimizable functions from files in the module root.
 
     Args:
@@ -478,9 +479,8 @@ def get_all_files_and_functions(
         Dictionary mapping file paths to lists of FunctionToOptimize.
 
     """
-    functions: dict[str, list[FunctionToOptimize]] = {}
+    functions: dict[Path, list[FunctionToOptimize]] = {}
     for file_path in get_files_for_language(module_root_path, ignore_paths, language):
-        # Find all the functions in the file
         functions.update(find_all_functions_in_file(file_path).items())
     # Randomize the order of the files to optimize to avoid optimizing the same file in the same order every time.
     # Helpful if an optimize-all run is stuck and we restart it.
@@ -785,7 +785,7 @@ def filter_functions(
     disable_logs: bool = False,
 ) -> tuple[dict[Path, list[FunctionToOptimize]], int]:
     resolved_project_root = project_root.resolve()
-    filtered_modified_functions: dict[str, list[FunctionToOptimize]] = {}
+    filtered_modified_functions: dict[Path, list[FunctionToOptimize]] = {}
     blocklist_funcs = get_blocklisted_functions()
     logger.debug(f"Blocklisted functions: {blocklist_funcs}")
     # Remove any function that we don't want to optimize
@@ -892,7 +892,7 @@ def is_test_file(file_path_normalized: str) -> bool:
                 functions_tmp.append(function)
             _functions = functions_tmp
 
-        filtered_modified_functions[file_path] = _functions
+        filtered_modified_functions[file_path_path] = _functions
         functions_count += len(_functions)
 
     if not disable_logs:
@@ -913,7 +913,7 @@ def is_test_file(file_path_normalized: str) -> bool:
         if len(tree.children) > 0:
             console.print(tree)
             console.rule()
-    return {Path(k): v for k, v in filtered_modified_functions.items() if v}, functions_count
+    return {k: v for k, v in filtered_modified_functions.items() if v}, functions_count
 
 
 def filter_files_optimized(file_path: Path, tests_root: Path, ignore_paths: list[Path], module_root: Path) -> bool:

codeflash/languages/python/support.py

@@ -134,14 +134,11 @@ def discover_functions(
         tree = cst.parse_module(source)
 
         wrapper = cst.metadata.MetadataWrapper(tree)
-        function_visitor = FunctionVisitor(file_path=str(file_path))
+        function_visitor = FunctionVisitor(file_path=file_path)
         wrapper.visit(function_visitor)
 
         functions: list[FunctionToOptimize] = []
         for func in function_visitor.functions:
-            if not isinstance(func, FunctionToOptimize):
-                continue
-
             if not criteria.include_async and func.is_async:
                 continue
 

codeflash/lsp/beta.py

@@ -114,7 +114,7 @@ def get_functions_in_commit(params: OptimizableFunctionsInCommitParams) -> dict[
     return {"functions": file_to_qualified_names, "status": "success"}
 
 
-def _group_functions_by_file(functions: dict[str, list[FunctionToOptimize]]) -> dict[str, list[str]]:
+def _group_functions_by_file(functions: dict[Path, list[FunctionToOptimize]]) -> dict[str, list[str]]:
     file_to_funcs_to_optimize, _ = filter_functions(
         modified_functions=functions,
         tests_root=server.optimizer.test_cfg.tests_root,