Check out previous branch rather than default branch on cleanup

Author: Jackenmen

cherry_picker/cherry_picker.py

@@ -38,6 +38,9 @@
     CHECKING_OUT_DEFAULT_BRANCH
     CHECKED_OUT_DEFAULT_BRANCH
 
+    CHECKING_OUT_PREVIOUS_BRANCH
+    CHECKED_OUT_PREVIOUS_BRANCH
+
     PUSHING_TO_REMOTE
     PUSHED_TO_REMOTE
     PUSHING_TO_REMOTE_FAILED
@@ -134,6 +137,11 @@ def set_paused_state(self):
             save_cfg_vals_to_git_cfg(config_path=self.chosen_config_path)
         set_state(WORKFLOW_STATES.BACKPORT_PAUSED)
 
+    def remember_previous_branch(self):
+        """Save the current branch into Git config to be able to get back to it later."""
+        current_branch = get_current_branch()
+        save_cfg_vals_to_git_cfg(previous_branch=current_branch)
+
     @property
     def upstream(self):
         """Get the remote name to use for upstream branches
@@ -180,24 +188,27 @@ def run_cmd(self, cmd):
         output = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
         return output.decode("utf-8")
 
-    def checkout_branch(self, branch_name):
+    def checkout_branch(self, branch_name, *, create_branch=False):
         """git checkout -b <branch_name>"""
-        cmd = [
-            "git",
-            "checkout",
-            "-b",
-            self.get_cherry_pick_branch(branch_name),
-            f"{self.upstream}/{branch_name}",
-        ]
+        if create_branch:
+            cmd = [
+                "git",
+                "checkout",
+                "-b",
+                branch_name,
+                f"{self.upstream}/{branch_name}",
+            ]
+        else:
+            cmd = ["git", "checkout", branch_name]
         try:
             self.run_cmd(cmd)
         except subprocess.CalledProcessError as err:
             click.echo(
-                f"Error checking out the branch {self.get_cherry_pick_branch(branch_name)}."
+                f"Error checking out the branch {branch_name}."
             )
             click.echo(err.output)
             raise BranchCheckoutException(
-                f"Error checking out the branch {self.get_cherry_pick_branch(branch_name)}."
+                f"Error checking out the branch {branch_name}."
             )
 
     def get_commit_message(self, commit_sha):
@@ -221,11 +232,23 @@ def checkout_default_branch(self):
         """git checkout default branch"""
         set_state(WORKFLOW_STATES.CHECKING_OUT_DEFAULT_BRANCH)
 
-        cmd = "git", "checkout", self.config["default_branch"]
-        self.run_cmd(cmd)
+        self.checkout_branch(self.config["default_branch"])
 
         set_state(WORKFLOW_STATES.CHECKED_OUT_DEFAULT_BRANCH)
 
+    def checkout_previous_branch(self):
+        """git checkout previous branch"""
+        set_state(WORKFLOW_STATES.CHECKING_OUT_PREVIOUS_BRANCH)
+
+        previous_branch = load_val_from_git_cfg("previous_branch")
+        if previous_branch is None:
+            self.checkout_default_branch()
+            return
+
+        self.checkout_branch(previous_branch)
+
+        set_state(WORKFLOW_STATES.CHECKED_OUT_PREVIOUS_BRANCH)
+
     def status(self):
         """
         git status
@@ -357,7 +380,12 @@ def cleanup_branch(self, branch):
         Switch to the default branch before that.
         """
         set_state(WORKFLOW_STATES.REMOVING_BACKPORT_BRANCH)
-        self.checkout_default_branch()
+        try:
+            self.checkout_previous_branch()
+        except BranchCheckoutException:
+            click.echo(f"branch {branch} NOT deleted.")
+            set_state(WORKFLOW_STATES.REMOVING_BACKPORT_BRANCH_FAILED)
+            return
         try:
             self.delete_branch(branch)
         except subprocess.CalledProcessError:
@@ -372,14 +400,15 @@ def backport(self):
             raise click.UsageError("At least one branch must be specified.")
         set_state(WORKFLOW_STATES.BACKPORT_STARTING)
         self.fetch_upstream()
+        self.remember_previous_branch()
 
         set_state(WORKFLOW_STATES.BACKPORT_LOOPING)
         for maint_branch in self.sorted_branches:
             set_state(WORKFLOW_STATES.BACKPORT_LOOP_START)
             click.echo(f"Now backporting '{self.commit_sha1}' into '{maint_branch}'")
 
             cherry_pick_branch = self.get_cherry_pick_branch(maint_branch)
-            self.checkout_branch(maint_branch)
+            self.checkout_branch(cherry_pick_branch, create_branch=True)
             commit_message = ""
             try:
                 self.cherry_pick()
@@ -413,6 +442,7 @@ def backport(self):
                     self.set_paused_state()
                     return  # to preserve the correct state
             set_state(WORKFLOW_STATES.BACKPORT_LOOP_END)
+        reset_stored_previous_branch()
         reset_state()
 
     def abort_cherry_pick(self):
@@ -434,6 +464,7 @@ def abort_cherry_pick(self):
         if get_current_branch().startswith("backport-"):
             self.cleanup_branch(get_current_branch())
 
+        reset_stored_previous_branch()
         reset_stored_config_ref()
         reset_state()
 
@@ -493,6 +524,7 @@ def continue_cherry_pick(self):
             )
             set_state(WORKFLOW_STATES.CONTINUATION_FAILED)
 
+        reset_stored_previous_branch()
         reset_stored_config_ref()
         reset_state()
 
@@ -822,6 +854,11 @@ def reset_stored_config_ref():
         """Config file pointer is not stored in Git config."""
 
 
+def reset_stored_previous_branch():
+    """Remove the previous branch information from Git config."""
+    wipe_cfg_vals_from_git_cfg("previous_branch")
+
+
 def reset_state():
     """Remove the progress state from Git config."""
     wipe_cfg_vals_from_git_cfg("state")