Fix stash pop and stash apply subcommands

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

Author: initialcommit-io

src/git_sim/commands.py

@@ -366,11 +366,15 @@ def stash(
         default=None,
         help="The name of the file to stash changes for",
     ),
+    stash_index: str = typer.Argument(
+        default="0",
+        help="Stash index",
+    ),
 ):
     from git_sim.stash import Stash
 
     settings.hide_first_tag = True
-    scene = Stash(files=files, command=command)
+    scene = Stash(files=files, command=command, stash_index=stash_index)
     handle_animations(scene=scene)
 
 

src/git_sim/stash.py

@@ -1,3 +1,4 @@
+import re
 import sys
 from enum import Enum
 import manim as m
@@ -10,14 +11,21 @@
 
 
 class Stash(GitSimBaseCommand):
-    def __init__(self, files: List[str], command: StashSubCommand):
+    def __init__(self, files: List[str], command: StashSubCommand, stash_index: int):
         super().__init__()
         self.files = files
         self.no_files = True if not self.files else False
         self.command = command
         settings.hide_merged_branches = True
         self.n = self.n_default
 
+        self.stash_index = self.parse_stash_format(stash_index)
+        if self.stash_index is None:
+            print(
+                "git-sim error: specify stash index as either integer or stash@{i}"
+            )
+            sys.exit()
+
         try:
             self.selected_branches.append(self.repo.active_branch.name)
         except TypeError:
@@ -44,7 +52,7 @@ def __init__(self, files: List[str], command: StashSubCommand):
                 and not settings.quiet
             ):
                 print(
-                    "Files are not required in apply/pop subcommand. Ignoring the file list....."
+                    "Files are not required in apply/pop subcommand. Ignoring the file list..."
                 )
 
         self.cmd += f"{type(self).__name__.lower()} {self.command.value if self.command else ''} {' '.join(self.files) if not self.no_files else ''}"
@@ -147,17 +155,21 @@ def populate_zones(
         thirdColumnArrowMap={},
     ):
         if self.command in [StashSubCommand.POP, StashSubCommand.APPLY]:
-            for x in self.repo.index.diff(None):
-                thirdColumnFileNames.add(x.a_path)
-                firstColumnFileNames.add(x.a_path)
-                thirdColumnArrowMap[x.a_path] = m.Arrow(
+            try:
+                stashedFileNames = self.repo.git.stash("show", "--name-only", self.stash_index)
+                stashedFileNames = stashedFileNames.split("\n")
+            except:
+                print(f"git-sim error: No stash entry with index {self.stashIndex} exists in stash")
+                sys.exit()
+            for s in stashedFileNames:
+                thirdColumnFileNames.add(s)
+                firstColumnFileNames.add(s)
+                thirdColumnArrowMap[s] = m.Arrow(
                     stroke_width=3, color=self.fontColor
                 )
-
-            for y in self.repo.index.diff("HEAD"):
-                firstColumnFileNames.add(y.a_path)
-                thirdColumnFileNames.add(y.a_path)
-                thirdColumnArrowMap[y.a_path] = m.Arrow(
+                firstColumnFileNames.add(s)
+                thirdColumnFileNames.add(s)
+                thirdColumnArrowMap[s] = m.Arrow(
                     stroke_width=3, color=self.fontColor
                 )
 
@@ -179,3 +191,14 @@ def populate_zones(
                         secondColumnArrowMap[y.a_path] = m.Arrow(
                             stroke_width=3, color=self.fontColor
                         )
+
+    def parse_stash_format(self, s):
+        # Regular expression to match either a plain integer or stash@{integer}
+        match = re.match(r"^(?:stash@\{(\d+)\}|\b(\d+)\b)$", s)
+        if match:
+            # match.group(1) is the integer in the stash@{integer} format
+            # match.group(2) is the integer if it's just a plain number
+            # One of these groups will be None, the other will have our number as a string
+            number_str = match.group(1) or match.group(2)
+            return int(number_str)
+        return None