Adjust rebase branch reset when using --onto with <until>, when <until> is specified as a non-active branch

Signed-off-by: Jacob Stopak <jacob@initialcommit.io>

Author: initialcommit-io

src/git_sim/git_sim_base_command.py

@@ -955,7 +955,7 @@ def center_frame_on_commit(self, commit):
         else:
             self.camera.frame.move_to(self.drawnCommits[commit.hexsha].get_center())
 
-    def reset_head_branch(self, hexsha, shift=numpy.array([0.0, 0.0, 0.0])):
+    def reset_head_branch(self, hexsha, branch="HEAD", shift=numpy.array([0.0, 0.0, 0.0])):
         if not self.head_exists():
             return
 
@@ -968,7 +968,7 @@ def reset_head_branch(self, hexsha, shift=numpy.array([0.0, 0.0, 0.0])):
                         0,
                     )
                 ),
-                self.drawnRefs[self.repo.active_branch.name].animate.move_to(
+                self.drawnRefs[self.repo.active_branch.name if branch == "HEAD" else branch].animate.move_to(
                     (
                         self.drawnCommits[hexsha].get_center()[0] + shift[0],
                         self.drawnCommits[hexsha].get_center()[1] + 2 + shift[1],
@@ -984,7 +984,7 @@ def reset_head_branch(self, hexsha, shift=numpy.array([0.0, 0.0, 0.0])):
                     0,
                 )
             )
-            self.drawnRefs[self.repo.active_branch.name].move_to(
+            self.drawnRefs[self.repo.active_branch.name if branch == "HEAD" else branch].move_to(
                 (
                     self.drawnCommits[hexsha].get_center()[0] + shift[0],
                     self.drawnCommits[hexsha].get_center()[1] + 2 + shift[1],

src/git_sim/rebase.py

@@ -171,8 +171,8 @@ def construct(self):
 
         if self.onto and self.until:
             until_sha = self.get_commit(self.until).hexsha
-            if until_sha == self.repo.head.commit.hexsha:
-                self.reset_head_branch(rebased_sha_map[until_sha])
+            if self.until in [b.name for b in self.repo.branches]:
+                self.reset_head_branch(rebased_sha_map[until_sha], branch=self.until)
             else:
                 try:
                     self.reset_head(rebased_sha_map[until_sha])