From ec839e9e9671e365b75cdc270d6f3602981fb9bb Mon Sep 17 00:00:00 2001 From: Jesse Duffield Date: Mon, 17 Jul 2023 13:40:26 +1000 Subject: [PATCH] Associate branches with worktrees even when mid-rebase --- pkg/commands/git_commands/branch_loader.go | 28 +++------ .../git_commands/branch_loader_test.go | 2 +- pkg/commands/git_commands/worktree.go | 21 +++++++ pkg/commands/git_commands/worktree_loader.go | 58 +++++++++++++++++++ pkg/commands/models/branch.go | 2 - pkg/gui/context/branches_context.go | 1 + pkg/gui/controllers/branches_controller.go | 22 +++---- pkg/gui/presentation/branches.go | 7 ++- 8 files changed, 104 insertions(+), 37 deletions(-) diff --git a/pkg/commands/git_commands/branch_loader.go b/pkg/commands/git_commands/branch_loader.go index ab14d65d1..cb284f67f 100644 --- a/pkg/commands/git_commands/branch_loader.go +++ b/pkg/commands/git_commands/branch_loader.go @@ -2,7 +2,6 @@ package git_commands import ( "fmt" - "os" "regexp" "strings" @@ -118,11 +117,6 @@ outer: } func (self *BranchLoader) obtainBranches() []*models.Branch { - currentDir, err := os.Getwd() - if err != nil { - panic(err) - } - output, err := self.getRawBranches() if err != nil { panic(err) @@ -144,7 +138,7 @@ func (self *BranchLoader) obtainBranches() []*models.Branch { return nil, false } - return obtainBranch(split, currentDir), true + return obtainBranch(split), true }) } @@ -172,32 +166,28 @@ var branchFields = []string{ "upstream:track", "subject", fmt.Sprintf("objectname:short=%d", utils.COMMIT_HASH_SHORT_SIZE), - "worktreepath", } // Obtain branch information from parsed line output of getRawBranches() -func obtainBranch(split []string, currentDir string) *models.Branch { +func obtainBranch(split []string) *models.Branch { headMarker := split[0] fullName := split[1] upstreamName := split[2] track := split[3] subject := split[4] commitHash := split[5] - branchDir := split[6] - checkedOutByOtherWorktree := len(branchDir) > 0 && branchDir != currentDir name := strings.TrimPrefix(fullName, "heads/") pushables, pullables, gone := parseUpstreamInfo(upstreamName, track) return &models.Branch{ - Name: name, - Pushables: pushables, - Pullables: pullables, - UpstreamGone: gone, - Head: headMarker == "*", - Subject: subject, - CommitHash: commitHash, - CheckedOutByOtherWorktree: checkedOutByOtherWorktree, + Name: name, + Pushables: pushables, + Pullables: pullables, + UpstreamGone: gone, + Head: headMarker == "*", + Subject: subject, + CommitHash: commitHash, } } diff --git a/pkg/commands/git_commands/branch_loader_test.go b/pkg/commands/git_commands/branch_loader_test.go index 86ee1640d..f16dcf5f4 100644 --- a/pkg/commands/git_commands/branch_loader_test.go +++ b/pkg/commands/git_commands/branch_loader_test.go @@ -81,7 +81,7 @@ func TestObtainBranch(t *testing.T) { for _, s := range scenarios { t.Run(s.testName, func(t *testing.T) { - branch := obtainBranch(s.input, "current-dir") + branch := obtainBranch(s.input) assert.EqualValues(t, s.expectedBranch, branch) }) } diff --git a/pkg/commands/git_commands/worktree.go b/pkg/commands/git_commands/worktree.go index 28e95d8b6..5a5637957 100644 --- a/pkg/commands/git_commands/worktree.go +++ b/pkg/commands/git_commands/worktree.go @@ -6,6 +6,8 @@ import ( "io/fs" "log" "os" + + "github.com/jesseduffield/lazygit/pkg/commands/models" ) type WorktreeCommands struct { @@ -84,3 +86,22 @@ func (self *WorktreeCommands) IsWorktreePathMissing(path string) bool { func EqualPath(a string, b string) bool { return a == b } + +func WorktreeForBranch(branch *models.Branch, worktrees []*models.Worktree) (*models.Worktree, bool) { + for _, worktree := range worktrees { + if worktree.Branch == branch.Name { + return worktree, true + } + } + + return nil, false +} + +func CheckedOutByOtherWorktree(branch *models.Branch, worktrees []*models.Worktree) bool { + worktree, ok := WorktreeForBranch(branch, worktrees) + if !ok { + return false + } + + return !IsCurrentWorktree(worktree.Path) +} diff --git a/pkg/commands/git_commands/worktree_loader.go b/pkg/commands/git_commands/worktree_loader.go index 23ccc39d8..75b60d433 100644 --- a/pkg/commands/git_commands/worktree_loader.go +++ b/pkg/commands/git_commands/worktree_loader.go @@ -2,6 +2,7 @@ package git_commands import ( "os" + "path/filepath" "strings" "github.com/jesseduffield/lazygit/pkg/commands/models" @@ -76,9 +77,66 @@ func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) { } } + // Some worktrees are on a branch but are mid-rebase, and in those cases, + // `git worktree list` will not show the branch name. We can get the branch + // name from the `rebase-merge/head-name` file (if it exists) in the folder + // for the worktree in the parent repo's .git/worktrees folder. + for _, worktree := range worktrees { + // No point checking if we already have a branch name + if worktree.Branch != "" { + continue + } + + rebaseBranch, ok := rebaseBranch(worktree.Path) + if ok { + worktree.Branch = rebaseBranch + } + } + return worktrees, nil } +func rebaseBranch(worktreePath string) (string, bool) { + // need to find the actual path of the worktree in the .git dir + gitPath, ok := worktreeGitPath(worktreePath) + if !ok { + return "", false + } + + // now we look inside that git path for a file `rebase-merge/head-name` + // if it exists, we update the worktree to say that it has that for a head + headNameContents, err := os.ReadFile(filepath.Join(gitPath, "rebase-merge", "head-name")) + if err != nil { + return "", false + } + + headName := strings.TrimSpace(string(headNameContents)) + shortHeadName := strings.TrimPrefix(headName, "refs/heads/") + + return shortHeadName, true +} + +func worktreeGitPath(worktreePath string) (string, bool) { + // first we get the path of the worktree, then we look at the contents of the `.git` file in that path + // then we look for the line that says `gitdir: /path/to/.git/worktrees/` + // then we return that path + gitFileContents, err := os.ReadFile(filepath.Join(worktreePath, ".git")) + if err != nil { + return "", false + } + + gitDirLine := lo.Filter(strings.Split(string(gitFileContents), "\n"), func(line string, _ int) bool { + return strings.HasPrefix(line, "gitdir: ") + }) + + if len(gitDirLine) == 0 { + return "", false + } + + gitDir := strings.TrimPrefix(gitDirLine[0], "gitdir: ") + return gitDir, true +} + type pathWithIndexT struct { path string index int diff --git a/pkg/commands/models/branch.go b/pkg/commands/models/branch.go index d3c9eb6dc..b4dcc0a79 100644 --- a/pkg/commands/models/branch.go +++ b/pkg/commands/models/branch.go @@ -26,8 +26,6 @@ type Branch struct { Subject string // commit hash CommitHash string - - CheckedOutByOtherWorktree bool } func (b *Branch) FullRefName() string { diff --git a/pkg/gui/context/branches_context.go b/pkg/gui/context/branches_context.go index 497b3a2c4..85c45c64a 100644 --- a/pkg/gui/context/branches_context.go +++ b/pkg/gui/context/branches_context.go @@ -31,6 +31,7 @@ func NewBranchesContext(c *ContextCommon) *BranchesContext { c.Modes().Diffing.Ref, c.Tr, c.UserConfig, + c.Model().Worktrees, ) } diff --git a/pkg/gui/controllers/branches_controller.go b/pkg/gui/controllers/branches_controller.go index 3ead83f35..4ea8099b2 100644 --- a/pkg/gui/controllers/branches_controller.go +++ b/pkg/gui/controllers/branches_controller.go @@ -202,11 +202,9 @@ func (self *BranchesController) press(selectedBranch *models.Branch) error { return self.c.ErrorMsg(self.c.Tr.AlreadyCheckedOutBranch) } - if selectedBranch.CheckedOutByOtherWorktree { - worktreeForRef, ok := self.worktreeForBranch(selectedBranch) - if ok && !self.c.Git().Worktree.IsCurrentWorktree(worktreeForRef.Path) { - return self.promptToCheckoutWorktree(worktreeForRef) - } + worktreeForRef, ok := self.worktreeForBranch(selectedBranch) + if ok && !self.c.Git().Worktree.IsCurrentWorktree(worktreeForRef.Path) { + return self.promptToCheckoutWorktree(worktreeForRef) } self.c.LogAction(self.c.Tr.Actions.CheckoutBranch) @@ -214,13 +212,7 @@ func (self *BranchesController) press(selectedBranch *models.Branch) error { } func (self *BranchesController) worktreeForBranch(branch *models.Branch) (*models.Worktree, bool) { - for _, worktree := range self.c.Model().Worktrees { - if worktree.Branch == branch.Name { - return worktree, true - } - } - - return nil, false + return git_commands.WorktreeForBranch(branch, self.c.Model().Worktrees) } func (self *BranchesController) promptToCheckoutWorktree(worktree *models.Worktree) error { @@ -326,13 +318,17 @@ func (self *BranchesController) delete(branch *models.Branch) error { return self.c.ErrorMsg(self.c.Tr.CantDeleteCheckOutBranch) } - if branch.CheckedOutByOtherWorktree { + if self.checkedOutByOtherWorktree(branch) { return self.promptWorktreeBranchDelete(branch) } return self.deleteWithForce(branch, false) } +func (self *BranchesController) checkedOutByOtherWorktree(branch *models.Branch) bool { + return git_commands.CheckedOutByOtherWorktree(branch, self.c.Model().Worktrees) +} + func (self *BranchesController) promptWorktreeBranchDelete(selectedBranch *models.Branch) error { worktree, ok := self.worktreeForBranch(selectedBranch) if !ok { diff --git a/pkg/gui/presentation/branches.go b/pkg/gui/presentation/branches.go index 42cbdf22b..6ad4e498e 100644 --- a/pkg/gui/presentation/branches.go +++ b/pkg/gui/presentation/branches.go @@ -5,6 +5,7 @@ import ( "strings" "github.com/jesseduffield/generics/slices" + "github.com/jesseduffield/lazygit/pkg/commands/git_commands" "github.com/jesseduffield/lazygit/pkg/commands/models" "github.com/jesseduffield/lazygit/pkg/config" "github.com/jesseduffield/lazygit/pkg/gui/presentation/icons" @@ -23,10 +24,11 @@ func GetBranchListDisplayStrings( diffName string, tr *i18n.TranslationSet, userConfig *config.UserConfig, + worktrees []*models.Worktree, ) [][]string { return slices.Map(branches, func(branch *models.Branch) []string { diffed := branch.Name == diffName - return getBranchDisplayStrings(branch, fullDescription, diffed, tr, userConfig) + return getBranchDisplayStrings(branch, fullDescription, diffed, tr, userConfig, worktrees) }) } @@ -37,6 +39,7 @@ func getBranchDisplayStrings( diffed bool, tr *i18n.TranslationSet, userConfig *config.UserConfig, + worktrees []*models.Worktree, ) []string { displayName := b.Name if b.DisplayName != "" { @@ -50,7 +53,7 @@ func getBranchDisplayStrings( coloredName := nameTextStyle.Sprint(displayName) branchStatus := utils.WithPadding(ColoredBranchStatus(b, tr), 2, utils.AlignLeft) - if b.CheckedOutByOtherWorktree { + if git_commands.CheckedOutByOtherWorktree(b, worktrees) { worktreeIcon := lo.Ternary(icons.IsIconEnabled(), icons.LINKED_WORKTREE_ICON, fmt.Sprintf("(%s)", tr.LcWorktree)) coloredName = fmt.Sprintf("%s %s", coloredName, style.FgDefault.Sprint(worktreeIcon)) }