Improve rebase subcommand to ignore merge commits by default, flatten diverging branch structure, color rebase dotted arrows, randomly generate simulated commit hashes

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

Author: initialcommit-io

src/git_sim/git_sim_base_command.py

@@ -1,18 +1,20 @@
 import os
-import platform
-import shutil
-import stat
+import git
 import sys
+import stat
+import numpy
+import random
+import shutil
+import platform
 import tempfile
 
-import git
 import manim as m
-import numpy
-from git.exc import GitCommandError, InvalidGitRepositoryError
+
 from git.repo import Repo
+from git.exc import GitCommandError, InvalidGitRepositoryError
 
-from git_sim.enums import ColorByOptions, StyleOptions
 from git_sim.settings import settings
+from git_sim.enums import ColorByOptions, StyleOptions
 
 
 class GitSimBaseCommand(m.MovingCameraScene):
@@ -402,7 +404,7 @@ def get_nonparent_branch_names(self):
         exclude = []
         for b1 in branches:
             for b2 in branches:
-                if b1.name != b2.name:
+                if b1.name != b2.name and b1.commit != b2.commit:
                     if self.repo.is_ancestor(b1.commit, b2.commit):
                         exclude.append(b1.name)
         return [b for b in branches if b.name not in exclude]
@@ -1375,6 +1377,10 @@ def add_ref_to_drawn_refs_by_commit(self, hexsha, ref):
                 ref,
             ]
 
+    def generate_random_sha(self):
+        valid_chars = "0123456789abcdef"
+        return "".join(random.choices(valid_chars, k=6))
+
 
 class DottedLine(m.Line):
     def __init__(self, *args, dot_spacing=0.4, dot_kwargs={}, **kwargs):

src/git_sim/rebase.py

@@ -1,11 +1,13 @@
 import sys
 
 import git
-import manim as m
 import numpy
+import random
+
+import manim as m
 
-from git_sim.git_sim_base_command import GitSimBaseCommand
 from git_sim.settings import settings
+from git_sim.git_sim_base_command import GitSimBaseCommand, DottedLine
 
 
 class Rebase(GitSimBaseCommand):
@@ -33,6 +35,16 @@ def __init__(self, branch: str):
 
         self.cmd += f"{type(self).__name__.lower()} {self.branch}"
 
+        self.alt_colors = {
+            0: [m.BLUE_B, m.BLUE_E],
+            1: [m.PURPLE_B, m.PURPLE_E],
+            2: [m.RED_B, m.RED_E],
+            3: [m.GREEN_B, m.GREEN_E],
+            4: [m.MAROON_B, m.MAROON_E],
+            5: [m.GOLD_B, m.GOLD_E],
+            6: [m.TEAL_B, m.TEAL_E],
+        }
+
     def construct(self):
         if not settings.stdout and not settings.output_only_path and not settings.quiet:
             print(f"{settings.INFO_STRING} {self.cmd}")
@@ -65,9 +77,12 @@ def construct(self):
         branch_commit = self.get_commit(self.branch)
         self.parse_commits(branch_commit)
         head_commit = self.get_commit()
+        default_commits = {}
+        self.get_default_commits(self.get_commit(), default_commits)
+        default_commits = self.sort_and_flatten(default_commits)
 
         reached_base = False
-        for commit in self.get_default_commits():
+        for commit in default_commits:
             if commit != "dark" and self.branch in self.repo.git.branch(
                 "--contains", commit
             ):
@@ -78,27 +93,45 @@ def construct(self):
         self.center_frame_on_commit(branch_commit)
 
         to_rebase = []
-        i = 0
-        current = head_commit
-        while self.branch not in self.repo.git.branch("--contains", current):
-            to_rebase.append(current)
-            i += 1
-            if i >= self.n:
-                break
-            current = self.get_default_commits()[i]
+        for c in default_commits:
+            if self.branch not in self.repo.git.branch("--contains", c):
+                to_rebase.append(c)
 
         parent = branch_commit.hexsha
-
+        branch_counts = {}
+        rebased_shas = []
         for j, tr in enumerate(reversed(to_rebase)):
+            if len(tr.parents) > 1:
+                continue
             if not reached_base and j == 0:
                 message = "..."
             else:
                 message = tr.message
-            parent = self.setup_and_draw_parent(parent, message)
-            self.draw_arrow_between_commits(tr.hexsha, parent)
+            color_index = int(self.drawnCommits[tr.hexsha].get_center()[1] / -4) - 1
+            if color_index not in branch_counts:
+                branch_counts[color_index] = 0
+            branch_counts[color_index] += 1
+            commit_color = self.alt_colors[color_index % len(self.alt_colors)][1]
+            parent = self.setup_and_draw_parent(parent, message, color=commit_color)
+            rebased_shas.append(parent)
 
         self.recenter_frame()
         self.scale_frame()
+
+        branch_counts = {}
+        k = 0
+        for j, tr in enumerate(reversed(to_rebase)):
+            if len(tr.parents) > 1:
+                k += 1
+                continue
+            color_index = int(self.drawnCommits[tr.hexsha].get_center()[1] / -4) - 1
+            if color_index not in branch_counts:
+                branch_counts[color_index] = 0
+            branch_counts[color_index] += 1
+            commit_color = self.alt_colors[color_index % len(self.alt_colors)][1]
+            arrow_color = self.alt_colors[color_index % len(self.alt_colors)][1 if branch_counts[color_index] % 2 == 0 else 1]
+            self.draw_arrow_between_commits(tr.hexsha, rebased_shas[j - k], color=arrow_color)
+
         self.reset_head_branch(parent)
         self.color_by(offset=2 * len(to_rebase))
         self.show_command_as_title()
@@ -111,11 +144,12 @@ def setup_and_draw_parent(
         commitMessage="New commit",
         shift=numpy.array([0.0, 0.0, 0.0]),
         draw_arrow=True,
+        color=m.RED,
     ):
         circle = m.Circle(
-            stroke_color=m.RED,
+            stroke_color=color,
             stroke_width=self.commit_stroke_width,
-            fill_color=m.RED,
+            fill_color=color,
             fill_opacity=0.25,
         )
         circle.height = 1
@@ -139,15 +173,9 @@ def setup_and_draw_parent(
         length = numpy.linalg.norm(start - end) - (1.5 if start[1] == end[1] else 3)
         arrow.set_length(length)
 
-        sha = "".join(
-            chr(ord(letter) + 1)
-            if (
-                (chr(ord(letter) + 1).isalpha() and letter < "f")
-                or chr(ord(letter) + 1).isdigit()
-            )
-            else letter
-            for letter in child[:6]
-        )
+        sha = None
+        while not sha or sha in self.drawnCommits:
+            sha = self.generate_random_sha()
         commitId = m.Text(
             sha if commitMessage != "..." else "...",
             font=self.font,
@@ -190,3 +218,28 @@ def setup_and_draw_parent(
             self.toFadeOut.add(arrow)
 
         return sha
+
+    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 commit not in self.sort_and_flatten(default_commits):
+                default_commits[branch_index].append(commit)
+                for i, parent in enumerate(commit.parents):
+                    self.get_default_commits(parent, default_commits, branch_index + i)
+        return default_commits
+
+    def draw_arrow_between_commits(self, startsha, endsha, color):
+        start = self.drawnCommits[startsha].get_center()
+        end = self.drawnCommits[endsha].get_center()
+
+        arrow = DottedLine(
+            start, end, color=color, dot_kwargs={"color": color}
+        ).add_tip()
+        length = numpy.linalg.norm(start - end) - 1.65
+        arrow.set_length(length)
+        self.draw_arrow(True, arrow)
+
+    def sort_and_flatten(self, d):
+        sorted_values = [d[key] for key in sorted(d.keys(), reverse=True)]
+        return sum(sorted_values, [])

src/git_sim/settings.py

@@ -20,8 +20,8 @@ class Settings(BaseSettings):
     transparent_bg: bool = False
     logo: pathlib.Path = pathlib.Path(__file__).parent.resolve() / "logo.png"
     low_quality: bool = False
-    max_branches_per_commit: int = 1
-    max_tags_per_commit: int = 1
+    max_branches_per_commit: int = 2
+    max_tags_per_commit: int = 2
     media_dir: pathlib.Path = pathlib.Path().cwd()
     outro_bottom_text: str = "Learn more at initialcommit.com"
     outro_top_text: str = "Thanks for using Initial Commit!"