Add --amend option for commit subcommand

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

Author: initialcommit-io

git_sim/__main__.py

@@ -133,6 +133,11 @@ def main():
         type=str,
         default="New commit",
     )
+    commit.add_argument(
+        "--amend",
+        help="Amend the last commit message, must be used with the -m flag",
+        action="store_true",
+    )
 
     stash = subparsers.add_parser("stash", help="stash -h")
     stash.add_argument(

git_sim/git_sim_base_command.py

@@ -13,6 +13,7 @@ def __init__(self, scene):
 
         self.drawnCommits = {}
         self.drawnRefs = {}
+        self.drawnCommitIds = {}
         self.commits = []
         self.zoomOuts = 0
         self.toFadeOut = m.Group()
@@ -219,6 +220,9 @@ def draw_commit(
         )
         commitId.next_to(circle, m.UP)
 
+        if commit != "dark":
+            self.drawnCommitIds[commit.hexsha] = commitId
+
         message = m.Text(
             "\n".join(
                 commitMessage[j : j + 20] for j in range(0, len(commitMessage), 20)

git_sim/git_sim_commit.py

@@ -11,38 +11,69 @@ class GitSimCommit(GitSimBaseCommand):
     def __init__(self, scene):
         super().__init__(scene)
         self.maxrefs = 2
-        self.defaultNumCommits = 4
-        self.numCommits = 4
+        self.defaultNumCommits = 4 if not self.scene.args.amend else 5
+        self.numCommits = 4 if not self.scene.args.amend else 5
         self.hide_first_tag = True
 
         try:
             self.selected_branches.append(self.repo.active_branch.name)
         except TypeError:
             pass
 
+        if self.scene.args.amend and self.scene.args.message == "New commit":
+            print(
+                "git-sim error: The --amend flag must be used with the -m flag to specify the amended commit message."
+            )
+            sys.exit(1)
+
     def execute(self):
         print(
             "Simulating: git "
             + self.scene.args.subcommand
+            + (" --amend" if self.scene.args.amend else "")
             + ' -m "'
             + self.scene.args.message
             + '"'
         )
 
         self.show_intro()
         self.get_commits()
+
+        if self.scene.args.amend:
+            tree = self.repo.tree()
+            amended = git.Commit.create_from_tree(
+                self.repo,
+                tree,
+                self.scene.args.message,
+            )
+            self.commits[0] = amended
+
         self.parse_commits(self.commits[self.i])
         self.center_frame_on_commit(self.commits[0])
-        self.setup_and_draw_parent(self.commits[0], self.scene.args.message)
+
+        if not self.scene.args.amend:
+            self.setup_and_draw_parent(self.commits[0], self.scene.args.message)
+        else:
+            self.draw_ref(self.commits[0], self.drawnCommitIds[amended.hexsha])
+            self.draw_ref(
+                self.commits[0],
+                self.drawnRefs["HEAD"],
+                text=self.repo.active_branch.name,
+                color=m.GREEN,
+            )
+
         self.recenter_frame()
         self.scale_frame()
-        self.reset_head_branch("abcdef")
-        self.vsplit_frame()
-        self.setup_and_draw_zones(
-            first_column_name="Working directory",
-            second_column_name="Staging area",
-            third_column_name="New commit",
-        )
+
+        if not self.scene.args.amend:
+            self.reset_head_branch("abcdef")
+            self.vsplit_frame()
+            self.setup_and_draw_zones(
+                first_column_name="Working directory",
+                second_column_name="Staging area",
+                third_column_name="New commit",
+            )
+
         self.fadeout()
         self.show_outro()