Add --onto flag to rebase subcommand to allow specification of parent of commit to rebase

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

Author: initialcommit-io

src/git_sim/commands.py

@@ -265,12 +265,23 @@ def rebase(
     ),
     rebase_merges: bool = typer.Option(
         False,
+        "--rebase-merges",
+        "-r",
         help="Preserve merge structure during rebase",
     ),
+    onto: bool = typer.Option(
+        False,
+        "--onto",
+        help="Rebase onto given branch instead of upstream",
+    ),
+    oldparent: str = typer.Argument(
+        None,
+        help="The parent of the commit to rebase (to be used with --onto)",
+    ),
 ):
     from git_sim.rebase import Rebase
 
-    scene = Rebase(branch=branch, rebase_merges=rebase_merges)
+    scene = Rebase(branch=branch, rebase_merges=rebase_merges, onto=onto, oldparent=oldparent)
     handle_animations(scene=scene)
 
 

src/git_sim/git_sim_base_command.py

@@ -11,7 +11,7 @@
 import manim as m
 
 from git.repo import Repo
-from git.exc import GitCommandError, InvalidGitRepositoryError
+from git.exc import GitCommandError, InvalidGitRepositoryError, BadName
 
 from git_sim.settings import settings
 from git_sim.enums import ColorByOptions, StyleOptions
@@ -106,7 +106,11 @@ def construct(self):
 
     def get_commit(self, sha_or_ref="HEAD"):
         if self.head_exists():
-            return self.repo.commit(sha_or_ref)
+            try:
+                return self.repo.commit(sha_or_ref)
+            except BadName:
+                print(f"git-sim error: {sha_or_ref} did not resolve to a valid Git object.")
+                sys.exit(1)
         return "dark"
 
     def get_default_commits(self):
@@ -1381,6 +1385,17 @@ def generate_random_sha(self):
         valid_chars = "0123456789abcdef"
         return "".join(random.choices(valid_chars, k=6))
 
+    def get_mainline_distance(self, sha_or_ref1, sha_or_ref2):
+        commit1 = self.get_commit(sha_or_ref1)
+        commit2 = self.get_commit(sha_or_ref2)
+        if not self.repo.is_ancestor(commit1, commit2):
+            print(f"git-sim error: specified sha/ref '{sha_or_ref1}' must be an ancestor of sha/ref '{sha_or_ref2}'.")
+            sys.exit(1)
+        d = 0
+        while self.get_commit(f"{commit2.hexsha}~{d}").hexsha != commit1.hexsha:
+            d += 1
+        return d
+
 
 class DottedLine(m.Line):
     def __init__(self, *args, dot_spacing=0.4, dot_kwargs={}, **kwargs):

src/git_sim/rebase.py

@@ -11,10 +11,12 @@
 
 
 class Rebase(GitSimBaseCommand):
-    def __init__(self, branch: str, rebase_merges: bool):
+    def __init__(self, branch: str, rebase_merges: bool, onto: bool, oldparent: str):
         super().__init__()
         self.branch = branch
         self.rebase_merges = rebase_merges
+        self.onto = onto
+        self.oldparent = oldparent
 
         try:
             git.repo.fun.rev_parse(self.repo, self.branch)
@@ -26,6 +28,14 @@ def __init__(self, branch: str, rebase_merges: bool):
             )
             sys.exit(1)
 
+        if self.onto:
+            if not self.oldparent:
+                print(
+                    "git-sim error: Please specify the parent of the commit to rebase ('oldparent')"
+                )
+                sys.exit(1)
+            self.n = self.get_mainline_distance(oldparent, "HEAD")
+
         if self.branch in [branch.name for branch in self.repo.heads]:
             self.selected_branches.append(self.branch)
 
@@ -34,7 +44,7 @@ def __init__(self, branch: str, rebase_merges: bool):
         except TypeError:
             pass
 
-        self.cmd += f"{type(self).__name__.lower()} {self.branch}"
+        self.cmd += f"{type(self).__name__.lower()}{' --rebase-merges' if self.rebase_merges else ''}{' --onto' if self.onto else ''} {self.branch}{' ' + self.oldparent if self.onto and self.oldparent else ''}"
 
         self.alt_colors = {
             0: [m.BLUE_B, m.BLUE_E],
@@ -81,25 +91,25 @@ def construct(self):
         self.get_default_commits(self.get_commit(), default_commits)
         flat_default_commits = self.sort_and_flatten(default_commits)
 
-        reached_base = False
-        merge_base = self.repo.git.merge_base(self.branch, self.repo.active_branch.name)
-        if merge_base in self.drawnCommits:
-            reached_base = True
-
         self.parse_commits(head_commit, shift=4 * m.DOWN)
         self.parse_all()
         self.center_frame_on_commit(branch_commit)
 
-        to_rebase = []
+        self.to_rebase = []
         for c in flat_default_commits:
             if self.branch not in self.repo.git.branch("--contains", c):
-                to_rebase.append(c)
+                self.to_rebase.append(c)
+
+        reached_base = False
+        merge_base = self.repo.git.merge_base(self.branch, self.repo.active_branch.name)
+        if merge_base in self.drawnCommits or (self.onto and self.to_rebase[-1].hexsha in self.drawnCommits):
+            reached_base = True
 
         parent = branch_commit.hexsha
         branch_counts = {}
         rebased_shas = []
         rebased_sha_map = {}
-        for j, tr in enumerate(reversed(to_rebase)):
+        for j, tr in enumerate(reversed(self.to_rebase)):
             if not self.rebase_merges:
                 if len(tr.parents) > 1:
                     continue
@@ -121,7 +131,7 @@ def construct(self):
 
         branch_counts = {}
         k = 0
-        for j, tr in enumerate(reversed(to_rebase)):
+        for j, tr in enumerate(reversed(self.to_rebase)):
             if not self.rebase_merges:
                 if len(tr.parents) > 1:
                     k += 1
@@ -138,7 +148,7 @@ def construct(self):
             self.reset_head_branch(rebased_sha_map[default_commits[0][0].hexsha])
         else:
             self.reset_head_branch(parent)
-        self.color_by(offset=2 * len(to_rebase))
+        self.color_by(offset=2 * len(self.to_rebase))
         self.show_command_as_title()
         self.fadeout()
         self.show_outro()
@@ -161,10 +171,21 @@ def setup_and_draw_parent(
             fill_opacity=0.25,
         )
         circle.height = 1
-        if self.rebase_merges and branch_index > 0:
+        side_offset = 0
+        num_branch_index_0_to_rebase = 0
+        for commit in default_commits[0]:
+            if commit in self.to_rebase:
+                num_branch_index_0_to_rebase += 1
+        if self.rebase_merges:
+            for bi in default_commits:
+                if bi > 0:
+                    if len(default_commits[bi]) >= num_branch_index_0_to_rebase:
+                        side_offset = len(default_commits[bi]) - num_branch_index_0_to_rebase + 1
+
+        if self.rebase_merges:
             circle.move_to(
                 self.drawnCommits[orig].get_center(),
-            ).shift(m.UP * 4 + (m.LEFT if settings.reverse else m.RIGHT) * len(default_commits[0]) * 2.5)
+            ).shift(m.UP * 4 + (m.LEFT if settings.reverse else m.RIGHT) * len(default_commits[0]) * 2.5 + (m.LEFT * side_offset if settings.reverse else m.RIGHT * side_offset) * 5)
         else:
             circle.next_to(
                 self.drawnCommits[child],
@@ -183,10 +204,10 @@ def setup_and_draw_parent(
             for p in self.get_commit(orig).parents:
                 if self.branch in self.repo.git.branch(
                     "--contains", p
-                ):
+                ) or p not in self.to_rebase:
                     continue
                 try:
-                    end = self.drawnCommits[p.hexsha].get_center() + m.UP * 4 + (m.LEFT if settings.reverse else m.RIGHT) * len(default_commits[0]) * 2.5
+                    end = self.drawnCommits[p.hexsha].get_center() + m.UP * 4 + (m.LEFT if settings.reverse else m.RIGHT) * len(default_commits[0]) * 2.5 + (m.LEFT * side_offset if settings.reverse else m.RIGHT * side_offset) * 5
                     arrow_start_ends.append((start, end))
                 except KeyError:
                     pass
@@ -257,6 +278,8 @@ def get_default_commits(self, commit, default_commits, branch_index=0):
         if branch_index not in default_commits:
             default_commits[branch_index] = []
         if len(default_commits[branch_index]) < self.n:
+            if self.onto and commit.hexsha == self.get_commit(self.oldparent).hexsha:
+                return default_commits
             if commit not in self.sort_and_flatten(default_commits) and self.branch not in self.repo.git.branch(
                 "--contains", commit
             ):