Add --set-upstream flag for push subcommand

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

Author: initialcommit-io

src/git_sim/commands.py

@@ -240,10 +240,15 @@ def push(
         default=None,
         help="The name of the branch to push",
     ),
+    set_upstream: bool = typer.Option(
+        False,
+        "--set-upstream",
+        help="Map the local branch to the specified upstream branch",
+    ),
 ):
     from git_sim.push import Push
 
-    scene = Push(remote=remote, branch=branch)
+    scene = Push(remote=remote, branch=branch, set_upstream=set_upstream)
     handle_animations(scene=scene)
 
 

src/git_sim/push.py

@@ -16,17 +16,18 @@
 
 
 class Push(GitSimBaseCommand):
-    def __init__(self, remote: str = None, branch: str = None):
+    def __init__(self, remote: str = None, branch: str = None, set_upstream: bool = False):
         super().__init__()
         self.remote = remote
         self.branch = branch
+        self.set_upstream = set_upstream
         settings.max_branches_per_commit = 2
 
         if self.remote and self.remote not in self.repo.remotes:
             print("git-sim error: no remote with name '" + self.remote + "'")
             sys.exit(1)
 
-        self.cmd += f"{type(self).__name__.lower()} {self.remote if self.remote else ''} {self.branch if self.branch else ''}"
+        self.cmd += f"{type(self).__name__.lower()} {'--set-upstream ' if self.set_upstream else ''}{self.remote if self.remote else ''} {self.branch if self.branch else ''}"
 
     def construct(self):
         if not settings.stdout and not settings.output_only_path and not settings.quiet: