Merge pull request #598 from codeflash-ai/git-branch-fixes

fix crash on unpushed branch push

Author: KRRT7

codeflash/code_utils/checkpoint.py

@@ -10,6 +10,8 @@
 
 from rich.prompt import Confirm
 
+from codeflash.cli_cmds.console import console
+
 if TYPE_CHECKING:
     import argparse
 
@@ -142,8 +144,11 @@ def ask_should_use_checkpoint_get_functions(args: argparse.Namespace) -> Optiona
         if previous_checkpoint_functions and Confirm.ask(
             "Previous Checkpoint detected from an incomplete optimization run, shall I continue the optimization from that point?",
             default=True,
+            console=console,
         ):
-            pass
+            console.rule()
         else:
             previous_checkpoint_functions = None
+
+    console.rule()
     return previous_checkpoint_functions

codeflash/code_utils/git_utils.py

@@ -22,7 +22,9 @@
     from git import Repo
 
 
-def get_git_diff(repo_directory: Path = Path.cwd(), uncommitted_changes: bool = False) -> dict[str, list[int]]:  # noqa: B008, FBT001, FBT002
+def get_git_diff(repo_directory: Path | None = None, *, uncommitted_changes: bool = False) -> dict[str, list[int]]:
+    if repo_directory is None:
+        repo_directory = Path.cwd()
     repository = git.Repo(repo_directory, search_parent_directories=True)
     commit = repository.head.commit
     if uncommitted_changes:
@@ -117,30 +119,31 @@ def confirm_proceeding_with_no_git_repo() -> str | bool:
     return True
 
 
-def check_and_push_branch(repo: git.Repo, git_remote: str | None = "origin", wait_for_push: bool = False) -> bool:  # noqa: FBT001, FBT002
-    current_branch = repo.active_branch.name
+def check_and_push_branch(repo: git.Repo, git_remote: str | None = "origin", *, wait_for_push: bool = False) -> bool:
+    current_branch = repo.active_branch
+    current_branch_name = current_branch.name
     remote = repo.remote(name=git_remote)
 
     # Check if the branch is pushed
-    if f"{git_remote}/{current_branch}" not in repo.refs:
-        logger.warning(f"⚠️ The branch '{current_branch}' is not pushed to the remote repository.")
+    if f"{git_remote}/{current_branch_name}" not in repo.refs:
+        logger.warning(f"⚠️ The branch '{current_branch_name}' is not pushed to the remote repository.")
         if not sys.__stdin__.isatty():
             logger.warning("Non-interactive shell detected. Branch will not be pushed.")
             return False
         if sys.__stdin__.isatty() and Confirm.ask(
             f"⚡️ In order for me to create PRs, your current branch needs to be pushed. Do you want to push "
-            f"the branch '{current_branch}' to the remote repository?",
+            f"the branch '{current_branch_name}' to the remote repository?",
             default=False,
         ):
             remote.push(current_branch)
-            logger.info(f"⬆️ Branch '{current_branch}' has been pushed to {git_remote}.")
+            logger.info(f"⬆️ Branch '{current_branch_name}' has been pushed to {git_remote}.")
             if wait_for_push:
                 time.sleep(3)  # adding this to give time for the push to register with GitHub,
                 # so that our modifications to it are not rejected
             return True
-        logger.info(f"🔘 Branch '{current_branch}' has not been pushed to {git_remote}.")
+        logger.info(f"🔘 Branch '{current_branch_name}' has not been pushed to {git_remote}.")
         return False
-    logger.debug(f"The branch '{current_branch}' is present in the remote repository.")
+    logger.debug(f"The branch '{current_branch_name}' is present in the remote repository.")
     return True
 
 

codeflash/discovery/functions_to_optimize.py

@@ -13,6 +13,7 @@
 import git
 import libcst as cst
 from pydantic.dataclasses import dataclass
+from rich.tree import Tree
 
 from codeflash.api.cfapi import get_blocklisted_functions, is_function_being_optimized_again
 from codeflash.cli_cmds.console import DEBUG_MODE, console, logger
@@ -37,6 +38,7 @@
 
     from codeflash.models.models import CodeOptimizationContext
     from codeflash.verification.verification_utils import TestConfig
+from rich.text import Text
 
 
 @dataclass(frozen=True)
@@ -594,20 +596,22 @@ def filter_functions(
 
     if not disable_logs:
         log_info = {
-            f"{test_functions_removed_count} test function{'s' if test_functions_removed_count != 1 else ''}": test_functions_removed_count,
-            f"{site_packages_removed_count} site-package function{'s' if site_packages_removed_count != 1 else ''}": site_packages_removed_count,
-            f"{malformed_paths_count} non-importable file path{'s' if malformed_paths_count != 1 else ''}": malformed_paths_count,
-            f"{non_modules_removed_count} function{'s' if non_modules_removed_count != 1 else ''} outside module-root": non_modules_removed_count,
-            f"{ignore_paths_removed_count} file{'s' if ignore_paths_removed_count != 1 else ''} from ignored paths": ignore_paths_removed_count,
-            f"{submodule_ignored_paths_count} file{'s' if submodule_ignored_paths_count != 1 else ''} from ignored submodules": submodule_ignored_paths_count,
-            f"{blocklist_funcs_removed_count} function{'s' if blocklist_funcs_removed_count != 1 else ''} as previously optimized": blocklist_funcs_removed_count,
-            f"{previous_checkpoint_functions_removed_count} function{'s' if previous_checkpoint_functions_removed_count != 1 else ''} skipped from checkpoint": previous_checkpoint_functions_removed_count,
+            "Test functions removed": (test_functions_removed_count, "yellow"),
+            "Site-package functions removed": (site_packages_removed_count, "magenta"),
+            "Non-importable file paths": (malformed_paths_count, "red"),
+            "Functions outside module-root": (non_modules_removed_count, "cyan"),
+            "Files from ignored paths": (ignore_paths_removed_count, "blue"),
+            "Files from ignored submodules": (submodule_ignored_paths_count, "bright_black"),
+            "Blocklisted functions removed": (blocklist_funcs_removed_count, "bright_red"),
+            "Functions skipped from checkpoint": (previous_checkpoint_functions_removed_count, "green"),
         }
-        log_string = "\n".join([k for k, v in log_info.items() if v > 0])
-        if log_string:
-            logger.info(f"Ignoring: {log_string}")
+        tree = Tree(Text("Ignored functions and files", style="bold"))
+        for label, (count, color) in log_info.items():
+            if count > 0:
+                tree.add(Text(f"{label}: {count}", style=color))
+        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
 
 

tests/test_git_utils.py

@@ -74,7 +74,7 @@ def test_check_and_push_branch(self, mock_confirm, mock_isatty, mock_repo):
         mock_origin.push.return_value = None
 
         assert check_and_push_branch(mock_repo_instance)
-        mock_origin.push.assert_called_once_with("test-branch")
+        mock_origin.push.assert_called_once_with(mock_repo_instance.active_branch)
         mock_origin.push.reset_mock()
 
         # Test when branch is already pushed