experimental changes

Author: mohammedahmed18

codeflash/cli_cmds/cli.py

@@ -155,6 +155,7 @@ def process_pyproject_config(args: Namespace) -> Namespace:
         "disable_telemetry",
         "disable_imports_sorting",
         "git_remote",
+        "exp_git_remote",
         "override_fixtures",
     ]
     for key in supported_keys:
@@ -234,7 +235,7 @@ def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
                 "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):
+        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:

codeflash/code_utils/config_parser.py

@@ -70,7 +70,7 @@ def parse_config_file(
     # default values:
     path_keys = ["module-root", "tests-root", "benchmarks-root"]
     path_list_keys = ["ignore-paths"]
-    str_keys = {"pytest-cmd": "pytest", "git-remote": "origin"}
+    str_keys = {"pytest-cmd": "pytest", "git-remote": "origin", "exp-git-remote": "exp-origin"}
     bool_keys = {
         "override-fixtures": False,
         "disable-telemetry": False,

codeflash/code_utils/git_utils.py

@@ -117,12 +117,12 @@ def confirm_proceeding_with_no_git_repo() -> str | bool:
     return True
 
 
-def check_and_push_branch(repo: git.Repo, wait_for_push: bool = False) -> bool:  # noqa: FBT001, FBT002
+def check_and_push_branch(repo: git.Repo, git_remote: str, wait_for_push: bool = False) -> bool:  # noqa: FBT001, FBT002
     current_branch = repo.active_branch.name
-    origin = repo.remote(name="origin")
+    remote = repo.remote(name=git_remote)
 
     # Check if the branch is pushed
-    if f"origin/{current_branch}" not in repo.refs:
+    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 not sys.__stdin__.isatty():
             logger.warning("Non-interactive shell detected. Branch will not be pushed.")
@@ -132,13 +132,13 @@ def check_and_push_branch(repo: git.Repo, wait_for_push: bool = False) -> bool:
             f"the branch '{current_branch}' to the remote repository?",
             default=False,
         ):
-            origin.push(current_branch)
-            logger.info(f"⬆️ Branch '{current_branch}' has been pushed to origin.")
+            remote.push(current_branch)
+            logger.info(f"⬆️ Branch '{current_branch}' 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 origin.")
+        logger.info(f"🔘 Branch '{current_branch}' has not been pushed to {git_remote}.")
         return False
     logger.debug(f"The branch '{current_branch}' is present in the remote repository.")
     return True

codeflash/optimization/function_optimizer.py

@@ -1041,7 +1041,7 @@ def find_and_process_best_optimization(
                         if self.experiment_id
                         else self.function_trace_id,
                         coverage_message=coverage_message,
-                        git_remote=self.args.git_remote,
+                        git_remote=self.args.git_remote if exp_type == "EXP0" else self.args.exp_git_remote,
                     )
                     if (
                         self.args.all

codeflash/result/create_pr.py

@@ -185,7 +185,7 @@ def check_create_pr(
         owner, repo = get_repo_owner_and_name(git_repo, git_remote)
         logger.info(f"Pushing to {git_remote} - Owner: {owner}, Repo: {repo}")
         console.rule()
-        if not check_and_push_branch(git_repo, wait_for_push=True):
+        if not check_and_push_branch(git_repo, git_remote, wait_for_push=True):
             logger.warning("⏭️ Branch is not pushed, skipping PR creation...")
             return
         relative_path = explanation.file_path.relative_to(git_root_dir()).as_posix()

tests/test_git_utils.py

@@ -73,13 +73,13 @@ def test_check_and_push_branch(self, mock_confirm, mock_isatty, mock_repo):
         mock_origin = mock_repo_instance.remote.return_value
         mock_origin.push.return_value = None
 
-        assert check_and_push_branch(mock_repo_instance)
+        assert check_and_push_branch(mock_repo_instance, git_remote="origin")
         mock_origin.push.assert_called_once_with("test-branch")
         mock_origin.push.reset_mock()
 
         # Test when branch is already pushed
         mock_repo_instance.refs = [f"origin/{mock_repo_instance.active_branch.name}"]
-        assert check_and_push_branch(mock_repo_instance)
+        assert check_and_push_branch(mock_repo_instance, git_remote="origin")
         mock_origin.push.assert_not_called()
         mock_origin.push.reset_mock()
 
@@ -93,7 +93,7 @@ def test_check_and_push_branch_non_tty(self, mock_isatty, mock_repo):
         mock_origin = mock_repo_instance.remote.return_value
         mock_origin.push.return_value = None
 
-        assert not check_and_push_branch(mock_repo_instance)
+        assert not check_and_push_branch(mock_repo_instance, git_remote="origin")
         mock_origin.push.assert_not_called()
         mock_origin.push.reset_mock()