From e29f09a22876bbf809f540c69006f57cb25eb7f4 Mon Sep 17 00:00:00 2001
From: Alexsaphir <git@alexsaphir.com>
Date: Wed, 8 May 2024 03:09:06 +0200
Subject: [PATCH] feat: allow diff to take a branch name as origin using
 worktree

---
 flux_local/git_repo.py  | 10 +++++++---
 flux_local/tool/diff.py | 11 ++++++++++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/flux_local/git_repo.py b/flux_local/git_repo.py
index 2ebf5881..a33a7742 100644
--- a/flux_local/git_repo.py
+++ b/flux_local/git_repo.py
@@ -772,17 +772,21 @@ async def update_kustomization(cluster: Cluster) -> None:
 
 
 @contextlib.contextmanager
-def create_worktree(repo: git.repo.Repo) -> Generator[Path, None, None]:
+def create_worktree(repo: git.repo.Repo, existing_branch: str | None = None) -> Generator[Path, None, None]:
     """Create a ContextManager for a new git worktree in the current repo.
 
     This is used to get a fork of the current repo without any local changes
     in order to produce a diff.
+    Specifying existing_branch allows  to compare the current state with the exisiting branch.
     """
     orig = os.getcwd()
     with tempfile.TemporaryDirectory() as tmp_dir:
         _LOGGER.debug("Creating worktree in %s", tmp_dir)
-        # Add --detach to avoid creating a branch since we will not make modifications
-        repo.git.worktree("add", "--detach", str(tmp_dir))
+        if existing_branch is None:
+            # Add --detach to avoid creating a branch since we will not make modifications
+            repo.git.worktree("add", "--detach", str(tmp_dir))
+        else:
+            repo.git.worktree("add", str(tmp_dir), existing_branch)
         os.chdir(tmp_dir)
         yield Path(tmp_dir)
     _LOGGER.debug("Restoring to %s", orig)
diff --git a/flux_local/tool/diff.py b/flux_local/tool/diff.py
index acdf8d78..e1e84dc7 100644
--- a/flux_local/tool/diff.py
+++ b/flux_local/tool/diff.py
@@ -198,6 +198,15 @@ def add_diff_flags(args: ArgumentParser) -> None:
         default=None,
         nargs="?",
     )
+
+    args.add_argument(
+        "--branch-orig",
+        help="Branch to compare against using worktree",
+        type=str,
+        default=None,
+        nargs="?",
+    )
+
     args.add_argument(
         "--strip-attrs",
         help="Labels or annotations to strip from the diff",
@@ -225,7 +234,7 @@ def create_diff_path(
         yield git_repo.PathSelector(path_orig, sources=kwargs.get("sources"))
         return
 
-    with git_repo.create_worktree(selector.repo) as worktree:
+    with git_repo.create_worktree(selector.repo, existing_branch=kwargs.get("branch_orig")) as worktree:
         yield git_repo.PathSelector(pathlib.Path(worktree) / selector.relative_path)