Merge remote-tracking branch 'origin/main' into cf-842

Author: aseembits93

codeflash/cli_cmds/cli.py

@@ -249,25 +249,29 @@ def project_root_from_module_root(module_root: Path, pyproject_file_path: Path)
 
 
 def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
-    if hasattr(args, "all"):
-        import git
-
-        from codeflash.code_utils.git_utils import check_and_push_branch, get_repo_owner_and_name
-        from codeflash.code_utils.github_utils import require_github_app_or_exit
-
-        # Ensure that the user can actually open PRs on the repo.
-        try:
-            git_repo = git.Repo(search_parent_directories=True)
-        except git.exc.InvalidGitRepositoryError:
-            logger.exception(
-                "I couldn't find a git repository in the current directory. "
-                "I need a git repository to run --all and open PRs for optimizations. Exiting..."
-            )
-            apologize_and_exit()
-        if not args.no_pr and not check_and_push_branch(git_repo, git_remote=args.git_remote):
-            exit_with_message("Branch is not pushed...", error_on_exit=True)
-        owner, repo = get_repo_owner_and_name(git_repo)
-        if not args.no_pr:
+    if hasattr(args, "all") or (hasattr(args, "file") and args.file):
+        no_pr = getattr(args, "no_pr", False)
+
+        if not no_pr:
+            import git
+
+            from codeflash.code_utils.git_utils import check_and_push_branch, get_repo_owner_and_name
+            from codeflash.code_utils.github_utils import require_github_app_or_exit
+
+            # Ensure that the user can actually open PRs on the repo.
+            try:
+                git_repo = git.Repo(search_parent_directories=True)
+            except git.exc.InvalidGitRepositoryError:
+                mode = "--all" if hasattr(args, "all") else "--file"
+                logger.exception(
+                    f"I couldn't find a git repository in the current directory. "
+                    f"I need a git repository to run {mode} and open PRs for optimizations. Exiting..."
+                )
+                apologize_and_exit()
+            git_remote = getattr(args, "git_remote", None)
+            if not check_and_push_branch(git_repo, git_remote=git_remote):
+                exit_with_message("Branch is not pushed...", error_on_exit=True)
+            owner, repo = get_repo_owner_and_name(git_repo)
             require_github_app_or_exit(owner, repo)
     if not hasattr(args, "all"):
         args.all = None

codeflash/cli_cmds/cmd_init.py

@@ -22,7 +22,7 @@
 from rich.table import Table
 from rich.text import Text
 
-from codeflash.api.cfapi import is_github_app_installed_on_repo
+from codeflash.api.cfapi import get_user_id, is_github_app_installed_on_repo
 from codeflash.cli_cmds.cli_common import apologize_and_exit
 from codeflash.cli_cmds.console import console, logger
 from codeflash.cli_cmds.extension import install_vscode_extension
@@ -1216,6 +1216,7 @@ def enter_api_key_and_save_to_rc() -> None:
         # On Windows, create a batch file in the user's home directory (not auto-run, just used to store api key)
         shell_rc_path.touch()
         click.echo(f"✅ Created {shell_rc_path}")
+    get_user_id(api_key=api_key)  # Used to verify whether the API key is valid.
     result = save_api_key_to_rc(api_key)
     if is_successful(result):
         click.echo(result.unwrap())
@@ -1373,7 +1374,7 @@ def ask_for_telemetry() -> bool:
     from rich.prompt import Confirm
 
     return Confirm.ask(
-        "⚡️ Would you like to enable telemetry to help us improve the Codeflash experience?",
+        "⚡️ Help us improve Codeflash by sharing anonymous usage data (e.g. errors encountered)?",
         default=True,
         show_default=True,
     )

codeflash/code_utils/config_parser.py

@@ -105,7 +105,7 @@ def parse_config_file(
         if lsp_mode:
             # don't fail in lsp mode if codeflash config is not found.
             return {}, config_file_path
-        msg = f"Could not find the 'codeflash' block in the config file {config_file_path}. Please run 'codeflash init' to create the config file."
+        msg = f"Could not find the 'codeflash' block in the config file {config_file_path}. Please run 'codeflash init' to add Codeflash config in the pyproject.toml config file."
         raise ValueError(msg) from e
     assert isinstance(config, dict)
 

codeflash/code_utils/formatter.py

@@ -97,6 +97,9 @@ def is_diff_line(line: str) -> bool:
 
 
 def format_generated_code(generated_test_source: str, formatter_cmds: list[str]) -> str:
+    formatter_name = formatter_cmds[0].lower() if formatter_cmds else "disabled"
+    if formatter_name == "disabled":  # nothing to do if no formatter provided
+        return re.sub(r"\n{2,}", "\n\n", generated_test_source)
     with tempfile.TemporaryDirectory() as test_dir_str:
         # try running formatter, if nothing changes (could be due to formatting failing or no actual formatting needed) return code with 2 or more newlines substituted with 2 newlines
         original_temp = Path(test_dir_str) / "original_temp.py"

codeflash/code_utils/git_worktree_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import configparser
 import subprocess
 import tempfile
 import time
@@ -18,14 +19,36 @@
 
 def create_worktree_snapshot_commit(worktree_dir: Path, commit_message: str) -> None:
     repository = git.Repo(worktree_dir, search_parent_directories=True)
-    with repository.config_writer() as cw:
+    username = None
+    no_username = False
+    email = None
+    no_email = False
+    with repository.config_reader(config_level="repository") as cr:
+        try:
+            username = cr.get("user", "name")
+        except (configparser.NoSectionError, configparser.NoOptionError):
+            no_username = True
+        try:
+            email = cr.get("user", "email")
+        except (configparser.NoSectionError, configparser.NoOptionError):
+            no_email = True
+    with repository.config_writer(config_level="repository") as cw:
         if not cw.has_option("user", "name"):
             cw.set_value("user", "name", "Codeflash Bot")
         if not cw.has_option("user", "email"):
             cw.set_value("user", "email", "bot@codeflash.ai")
 
     repository.git.add(".")
     repository.git.commit("-m", commit_message, "--no-verify")
+    with repository.config_writer(config_level="repository") as cw:
+        if username:
+            cw.set_value("user", "name", username)
+        elif no_username:
+            cw.remove_option("user", "name")
+        if email:
+            cw.set_value("user", "email", email)
+        elif no_email:
+            cw.remove_option("user", "email")
 
 
 def create_detached_worktree(module_root: Path) -> Optional[Path]:

codeflash/code_utils/instrument_existing_tests.py

@@ -684,27 +684,6 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
                     )
 
 
-def instrument_source_module_with_async_decorators(
-    source_path: Path, function_to_optimize: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
-) -> tuple[bool, str | None]:
-    if not function_to_optimize.is_async:
-        return False, None
-
-    try:
-        with source_path.open(encoding="utf8") as f:
-            source_code = f.read()
-
-        modified_code, decorator_added = add_async_decorator_to_function(source_code, function_to_optimize, mode)
-
-        if decorator_added:
-            return True, modified_code
-
-    except Exception:
-        return False, None
-    else:
-        return False, None
-
-
 def inject_async_profiling_into_existing_test(
     test_path: Path,
     call_positions: list[CodePosition],
@@ -1288,25 +1267,29 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
 
 
 def add_async_decorator_to_function(
-    source_code: str, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
-) -> tuple[str, bool]:
-    """Add async decorator to an async function definition.
+    source_path: Path, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
+) -> bool:
+    """Add async decorator to an async function definition and write back to file.
 
     Args:
     ----
-        source_code: The source code to modify.
+        source_path: Path to the source file to modify in-place.
         function: The FunctionToOptimize object representing the target async function.
         mode: The testing mode to determine which decorator to apply.
 
     Returns:
     -------
-        Tuple of (modified_source_code, was_decorator_added).
+        Boolean indicating whether the decorator was successfully added.
 
     """
     if not function.is_async:
-        return source_code, False
+        return False
 
     try:
+        # Read source code
+        with source_path.open(encoding="utf8") as f:
+            source_code = f.read()
+
         module = cst.parse_module(source_code)
 
         # Add the decorator to the function
@@ -1318,10 +1301,17 @@ def add_async_decorator_to_function(
             import_transformer = AsyncDecoratorImportAdder(mode)
             module = module.visit(import_transformer)
 
-        return sort_imports(code=module.code, float_to_top=True), decorator_transformer.added_decorator
+        modified_code = sort_imports(code=module.code, float_to_top=True)
     except Exception as e:
         logger.exception(f"Error adding async decorator to function {function.qualified_name}: {e}")
-        return source_code, False
+        return False
+    else:
+        if decorator_transformer.added_decorator:
+            with source_path.open("w", encoding="utf8") as f:
+                f.write(modified_code)
+            logger.debug(f"Applied async {mode.value} instrumentation to {source_path}")
+            return True
+        return False
 
 
 def create_instrumented_source_module_path(source_path: Path, temp_dir: Path) -> Path:

codeflash/context/unused_definition_remover.py

@@ -469,22 +469,32 @@ def remove_unused_definitions_by_function_names(code: str, qualified_function_na
         qualified_function_names: Set of function names to keep. For methods, use format 'classname.methodname'
 
     """
-    module = cst.parse_module(code)
-    # Collect all definitions (top level classes, variables or function)
-    definitions = collect_top_level_definitions(module)
+    try:
+        module = cst.parse_module(code)
+    except Exception as e:
+        logger.debug(f"Failed to parse code with libcst: {type(e).__name__}: {e}")
+        return code
 
-    # Collect dependencies between definitions using the visitor pattern
-    dependency_collector = DependencyCollector(definitions)
-    module.visit(dependency_collector)
+    try:
+        # Collect all definitions (top level classes, variables or function)
+        definitions = collect_top_level_definitions(module)
 
-    # Mark definitions used by specified functions, and their dependencies recursively
-    usage_marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
-    usage_marker.mark_used_definitions()
+        # Collect dependencies between definitions using the visitor pattern
+        dependency_collector = DependencyCollector(definitions)
+        module.visit(dependency_collector)
 
-    # Apply the recursive removal transformation
-    modified_module, _ = remove_unused_definitions_recursively(module, definitions)
+        # Mark definitions used by specified functions, and their dependencies recursively
+        usage_marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
+        usage_marker.mark_used_definitions()
 
-    return modified_module.code if modified_module else ""
+        # Apply the recursive removal transformation
+        modified_module, _ = remove_unused_definitions_recursively(module, definitions)
+
+        return modified_module.code if modified_module else ""  # noqa: TRY300
+    except Exception as e:
+        # If any other error occurs during processing, return the original code
+        logger.debug(f"Error processing code to remove unused definitions: {type(e).__name__}: {e}")
+        return code
 
 
 def print_definitions(definitions: dict[str, UsageInfo]) -> None:

codeflash/discovery/functions_to_optimize.py

@@ -201,7 +201,7 @@ def get_functions_to_optimize(
         elif file is not None:
             logger.info("!lsp|Finding all functions in the file '%s'…", file)
             console.rule()
-            functions = find_all_functions_in_file(file)
+            functions: dict[Path, list[FunctionToOptimize]] = 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:
@@ -224,8 +224,16 @@ def get_functions_to_optimize(
                 if found_function is None:
                     if is_lsp:
                         return functions, 0, None
+                    found = closest_matching_file_function_name(only_get_this_function, functions)
+                    if found is not None:
+                        file, found_function = found
+                        exit_with_message(
+                            f"Function {only_get_this_function} not found in file {file}\nor the function does not have a 'return' statement or is a property.\n"
+                            f"Did you mean {found_function.qualified_name} instead?"
+                        )
+
                     exit_with_message(
-                        f"Function {only_function_name} not found in file {file}\nor the function does not have a 'return' statement or is a property"
+                        f"Function {only_get_this_function} not found in file {file}\nor the function does not have a 'return' statement or is a property"
                     )
                 functions[file] = [found_function]
         else:
@@ -259,6 +267,76 @@ def get_functions_within_git_diff(uncommitted_changes: bool) -> dict[str, list[F
     return get_functions_within_lines(modified_lines)
 
 
+def closest_matching_file_function_name(
+    qualified_fn_to_find: str, found_fns: dict[Path, list[FunctionToOptimize]]
+) -> tuple[Path, FunctionToOptimize] | None:
+    """Find the closest matching function name using Levenshtein distance.
+
+    Args:
+        qualified_fn_to_find: Function name to find in format "Class.function" or "function"
+        found_fns: Dictionary of file paths to list of functions
+
+    Returns:
+        Tuple of (file_path, function) for closest match, or None if no matches found
+
+    """
+    min_distance = 4
+    closest_match = None
+    closest_file = None
+
+    qualified_fn_to_find_lower = qualified_fn_to_find.lower()
+
+    # Cache levenshtein_distance locally for improved lookup speed
+    _levenshtein = levenshtein_distance
+
+    for file_path, functions in found_fns.items():
+        for function in functions:
+            # Compare either full qualified name or just function name
+            fn_name = function.qualified_name.lower()
+            # If the absolute length difference is already >= min_distance, skip calculation
+            if abs(len(qualified_fn_to_find_lower) - len(fn_name)) >= min_distance:
+                continue
+            dist = _levenshtein(qualified_fn_to_find_lower, fn_name)
+
+            if dist < min_distance:
+                min_distance = dist
+                closest_match = function
+                closest_file = file_path
+
+    if closest_match is not None:
+        return closest_file, closest_match
+    return None
+
+
+def levenshtein_distance(s1: str, s2: str) -> int:
+    if len(s1) > len(s2):
+        s1, s2 = s2, s1
+    len1 = len(s1)
+    len2 = len(s2)
+    # Use a preallocated list instead of creating a new list every iteration
+    previous = list(range(len1 + 1))
+    current = [0] * (len1 + 1)
+
+    for index2 in range(len2):
+        char2 = s2[index2]
+        current[0] = index2 + 1
+        for index1 in range(len1):
+            char1 = s1[index1]
+            if char1 == char2:
+                current[index1 + 1] = previous[index1]
+            else:
+                # Fast min calculation without tuple construct
+                a = previous[index1]
+                b = previous[index1 + 1]
+                c = current[index1]
+                min_val = min(b, a)
+                min_val = min(c, min_val)
+                current[index1 + 1] = 1 + min_val
+        # Swap references instead of copying
+        previous, current = current, previous
+    return previous[len1]
+
+
 def get_functions_inside_a_commit(commit_hash: str) -> dict[str, list[FunctionToOptimize]]:
     modified_lines: dict[str, list[int]] = get_git_diff(only_this_commit=commit_hash)
     return get_functions_within_lines(modified_lines)
@@ -405,7 +483,10 @@ def is_git_repo(file_path: str) -> bool:
 def ignored_submodule_paths(module_root: str) -> list[str]:
     if is_git_repo(module_root):
         git_repo = git.Repo(module_root, search_parent_directories=True)
-        return [Path(git_repo.working_tree_dir, submodule.path).resolve() for submodule in git_repo.submodules]
+        try:
+            return [Path(git_repo.working_tree_dir, submodule.path).resolve() for submodule in git_repo.submodules]
+        except Exception as e:
+            logger.warning(f"Error getting submodule paths: {e}")
     return []