diff --git a/fix_checkout_branch.py b/fix_checkout_branch.py new file mode 100644 index 000000000..9b3187a33 --- /dev/null +++ b/fix_checkout_branch.py @@ -0,0 +1,100 @@ +def checkout_branch(self, branch_name: str | None, *, remote: bool = False, remote_name: str = "origin", create_if_missing: bool = True) -> CheckoutResult: + """Attempts to check out the branch in the following order: + - Check out the local branch by name + - Check out the remote branch if it's been fetched + - Creates a new branch from the current commit (with create=True) + + NOTE: if branch is already checked out this does nothing. + TIP: Use remote=True if you want to always try to checkout the branch from a remote + + Args: + ---- + branch_name (str): Name of the branch to checkout. + create_if_missing: If the branch doesn't exist, create one + remote: Checks out a branch from a Remote + tracks the Remote + force (bool): If True, force checkout by resetting the current branch to HEAD. + If False, raise an error if the branch is dirty. + + Raises: + ------ + GitCommandError: If there's an error with Git operations. + RuntimeError: If the branch is dirty and force is not set. + """ + if branch_name is None: + branch_name = self.default_branch + + try: + if self.is_branch_checked_out(branch_name): + if remote: + # If the branch is already checked out and we want to fetch it from the remote, reset --hard to the remote branch + logger.info(f"Branch {branch_name} is already checked out locally. Resetting to remote branch: {remote_name}/{branch_name}") + # TODO: would have to fetch the the remote branch first to retrieve latest changes + self.git_cli.git.reset("--hard", f"{remote_name}/{branch_name}") + return CheckoutResult.SUCCESS + else: + logger.info(f"Branch {branch_name} is already checked out! Skipping checkout_branch.") + return CheckoutResult.SUCCESS + + # Check if there are changes that need to be preserved + needs_stash = self.git_cli.is_dirty() + if needs_stash: + logger.info(f"Environment is dirty, stashing changes before checking out branch: {branch_name}.") + self.stash_push() + + try: + # If remote=True, create a local branch tracking the remote branch and checkout onto it + if remote: + res = self.fetch_remote(remote_name, refspec=f"{branch_name}:{branch_name}") + if res is FetchResult.SUCCESS: + self.git_cli.git.checkout(branch_name) + if needs_stash: + logger.info(f"Applying stashed changes after checkout to branch: {branch_name}.") + self.stash_pop() + return CheckoutResult.SUCCESS + if res is FetchResult.REFSPEC_NOT_FOUND: + logger.warning(f"Branch {branch_name} not found in remote {remote_name}. Unable to checkout remote branch.") + if needs_stash: + logger.info("Restoring stashed changes.") + self.stash_pop() + return CheckoutResult.NOT_FOUND + + # If the branch already exists, checkout onto it + if branch_name in self.git_cli.heads: + self.git_cli.heads[branch_name].checkout() + if needs_stash: + logger.info(f"Applying stashed changes after checkout to branch: {branch_name}.") + self.stash_pop() + return CheckoutResult.SUCCESS + + # If the branch does not exist and create_if_missing=True, create and checkout a new branch from the current commit + elif create_if_missing: + logger.info(f"Creating new branch {branch_name} from current commit: {self.git_cli.head.commit.hexsha}") + new_branch = self.git_cli.create_head(branch_name) + new_branch.checkout() + if needs_stash: + logger.info(f"Applying stashed changes after checkout to new branch: {branch_name}.") + self.stash_pop() + return CheckoutResult.SUCCESS + else: + if needs_stash: + logger.info("Restoring stashed changes.") + self.stash_pop() + return CheckoutResult.NOT_FOUND + + except Exception as e: + # If anything goes wrong, try to restore the stashed changes + if needs_stash: + try: + logger.info("An error occurred. Attempting to restore stashed changes.") + self.stash_pop() + except Exception as stash_error: + logger.exception(f"Failed to restore stashed changes: {stash_error}") + raise e + + except GitCommandError as e: + if "fatal: ambiguous argument" in e.stderr: + logger.warning(f"Branch {branch_name} was not found in remote {remote_name}. Unable to checkout.") + return CheckoutResult.NOT_FOUND + else: + logger.exception(f"Error with Git operations: {e}") + raise diff --git a/src/codegen/git/repo_operator/repo_operator.py b/src/codegen/git/repo_operator/repo_operator.py index edee45a18..07540bbc3 100644 --- a/src/codegen/git/repo_operator/repo_operator.py +++ b/src/codegen/git/repo_operator/repo_operator.py @@ -415,7 +415,19 @@ def checkout_branch(self, branch_name: str | None, *, remote: bool = False, remo logger.info(f"Branch {branch_name} is already checked out! Skipping checkout_branch.") return CheckoutResult.SUCCESS - if self.git_cli.is_dirty(): +<<<<<<< Updated upstream + # Check if we need to preserve changes + is_dirty = self.git_cli.is_dirty() + need_to_create_branch = create_if_missing and branch_name not in self.git_cli.heads + + # If we're creating a new branch and there are changes, stash them + stashed_changes = False + if is_dirty and need_to_create_branch: + logger.info("Environment is dirty and creating new branch. Stashing changes before checkout.") + self.stash_push() + stashed_changes = True + # Otherwise, if we're not creating a new branch and there are changes, discard them + elif is_dirty: logger.info(f"Environment is dirty, discarding changes before checking out branch: {branch_name}.") self.discard_changes() @@ -424,14 +436,26 @@ def checkout_branch(self, branch_name: str | None, *, remote: bool = False, remo res = self.fetch_remote(remote_name, refspec=f"{branch_name}:{branch_name}") if res is FetchResult.SUCCESS: self.git_cli.git.checkout(branch_name) + # Apply stashed changes if needed + if stashed_changes: + logger.info("Applying stashed changes after checkout.") + self.stash_pop() return CheckoutResult.SUCCESS if res is FetchResult.REFSPEC_NOT_FOUND: logger.warning(f"Branch {branch_name} not found in remote {remote_name}. Unable to checkout remote branch.") + # Apply stashed changes if needed + if stashed_changes: + logger.info("Applying stashed changes after failed checkout.") + self.stash_pop() return CheckoutResult.NOT_FOUND # If the branch already exists, checkout onto it if branch_name in self.git_cli.heads: self.git_cli.heads[branch_name].checkout() + # Apply stashed changes if needed + if stashed_changes: + logger.info("Applying stashed changes after checkout.") + self.stash_pop() return CheckoutResult.SUCCESS # If the branch does not exist and create_if_missing=True, create and checkout a new branch from the current commit @@ -439,11 +463,84 @@ def checkout_branch(self, branch_name: str | None, *, remote: bool = False, remo logger.info(f"Creating new branch {branch_name} from current commit: {self.git_cli.head.commit.hexsha}") new_branch = self.git_cli.create_head(branch_name) new_branch.checkout() + # Apply stashed changes if needed + if stashed_changes: + logger.info("Applying stashed changes after creating new branch.") + self.stash_pop() return CheckoutResult.SUCCESS else: + # Apply stashed changes if needed + if stashed_changes: + logger.info("Applying stashed changes after failed checkout.") + self.stash_pop() return CheckoutResult.NOT_FOUND +======= + # Check if there are changes that need to be preserved + needs_stash = self.git_cli.is_dirty() + if needs_stash: + logger.info(f"Environment is dirty, stashing changes before checking out branch: {branch_name}.") + self.stash_push() + + try: + # If remote=True, create a local branch tracking the remote branch and checkout onto it + if remote: + res = self.fetch_remote(remote_name, refspec=f"{branch_name}:{branch_name}") + if res is FetchResult.SUCCESS: + self.git_cli.git.checkout(branch_name) + if needs_stash: + logger.info(f"Applying stashed changes after checkout to branch: {branch_name}.") + self.stash_pop() + return CheckoutResult.SUCCESS + if res is FetchResult.REFSPEC_NOT_FOUND: + logger.warning(f"Branch {branch_name} not found in remote {remote_name}. Unable to checkout remote branch.") + if needs_stash: + logger.info("Restoring stashed changes.") + self.stash_pop() + return CheckoutResult.NOT_FOUND + + # If the branch already exists, checkout onto it + if branch_name in self.git_cli.heads: + self.git_cli.heads[branch_name].checkout() + if needs_stash: + logger.info(f"Applying stashed changes after checkout to branch: {branch_name}.") + self.stash_pop() + return CheckoutResult.SUCCESS + + # If the branch does not exist and create_if_missing=True, create and checkout a new branch from the current commit + elif create_if_missing: + logger.info(f"Creating new branch {branch_name} from current commit: {self.git_cli.head.commit.hexsha}") + new_branch = self.git_cli.create_head(branch_name) + new_branch.checkout() + if needs_stash: + logger.info(f"Applying stashed changes after checkout to new branch: {branch_name}.") + self.stash_pop() + return CheckoutResult.SUCCESS + else: + if needs_stash: + logger.info("Restoring stashed changes.") + self.stash_pop() + return CheckoutResult.NOT_FOUND + + except Exception as e: + # If anything goes wrong, try to restore the stashed changes + if needs_stash: + try: + logger.info("An error occurred. Attempting to restore stashed changes.") + self.stash_pop() + except Exception as stash_error: + logger.error(f"Failed to restore stashed changes: {stash_error}") + raise e +>>>>>>> Stashed changes except GitCommandError as e: + # Apply stashed changes if needed + if "stashed_changes" in locals() and stashed_changes: + logger.info("Applying stashed changes after error.") + try: + self.stash_pop() + except Exception as stash_error: + logger.exception(f"Failed to apply stashed changes: {stash_error}") + if "fatal: ambiguous argument" in e.stderr: logger.warning(f"Branch {branch_name} was not found in remote {remote_name}. Unable to checkout.") return CheckoutResult.NOT_FOUND