better upstream tracking and allow renaming a branch

7 files changed, 106 insertions, 58 deletions

diff --git a/pkg/gui/branches_panel.go b/pkg/gui/branches_panel.go
index caac8e218..b2e9b1525 100644
--- a/pkg/gui/branches_panel.go
+++ b/pkg/gui/branches_panel.go
@@ -41,9 +41,6 @@ func (gui *Gui) handleBranchSelect(g *gocui.Gui, v *gocui.View) error {
 	}
 	branch := gui.getSelectedBranch()
 	v.FocusPoint(0, gui.State.Panels.Branches.SelectedLine)
-	if err := gui.RenderSelectedBranchUpstreamDifferences(); err != nil {
-		return err
-	}
 
 	cmd := gui.OSCommand.ExecutableFromString(
 		gui.GitCommand.GetBranchGraphCmdStr(branch.Name),
@@ -54,24 +51,6 @@ func (gui *Gui) handleBranchSelect(g *gocui.Gui, v *gocui.View) error {
 	return nil
 }
 
-func (gui *Gui) RenderSelectedBranchUpstreamDifferences() error {
-	return gui.newTask("branches", func(stop chan struct{}) error {
-		branch := gui.getSelectedBranch()
-		branch.Pushables, branch.Pullables = gui.GitCommand.GetBranchUpstreamDifferenceCount(branch.Name)
-
-		select {
-		case <-stop:
-			return nil
-		default:
-		}
-
-		branchesView := gui.getBranchesView()
-		displayStrings := presentation.GetBranchListDisplayStrings(gui.State.Branches, gui.currentViewName() == "branches", gui.State.Panels.Branches.SelectedLine)
-		gui.renderDisplayStrings(branchesView, displayStrings)
-		return nil
-	})
-}
-
 // gui.refreshStatus is called at the end of this because that's when we can
 // be sure there is a state.Branches array to pick the current branch from
 func (gui *Gui) refreshBranches(g *gocui.Gui) error {
@@ -106,9 +85,8 @@ func (gui *Gui) renderLocalBranchesWithSelection() error {
 	branchesView := gui.getBranchesView()
 
 	gui.refreshSelectedLine(&gui.State.Panels.Branches.SelectedLine, len(gui.State.Branches))
-	if err := gui.RenderSelectedBranchUpstreamDifferences(); err != nil {
-		return err
-	}
+	displayStrings := presentation.GetBranchListDisplayStrings(gui.State.Branches, gui.State.ScreenMode != SCREEN_NORMAL)
+	gui.renderDisplayStrings(branchesView, displayStrings)
 	if gui.g.CurrentView() == branchesView {
 		if err := gui.handleBranchSelect(gui.g, branchesView); err != nil {
 			return err
@@ -375,7 +353,7 @@ func (gui *Gui) handleFastForward(g *gocui.Gui, v *gocui.View) error {
 			if err := gui.GitCommand.FastForward(branch.Name, remoteName, remoteBranchName); err != nil {
 				_ = gui.createErrorPanel(gui.g, err.Error())
 			}
-			_ = gui.RenderSelectedBranchUpstreamDifferences()
+			_ = gui.refreshBranches(gui.g)
 		}
 
 		_ = gui.closeConfirmationPrompt(gui.g, true)
@@ -407,7 +385,13 @@ func (gui *Gui) switchBranchesPanelContext(context string) error {
 
 	branchesView.TabIndex = contextTabIndexMap[context]
 
-	switch context {
+	return gui.refreshBranchesViewWithSelection()
+}
+
+func (gui *Gui) refreshBranchesViewWithSelection() error {
+	branchesView := gui.getBranchesView()
+
+	switch branchesView.Context {
 	case "local-branches":
 		return gui.renderLocalBranchesWithSelection()
 	case "remotes":
@@ -457,3 +441,41 @@ func (gui *Gui) onBranchesPanelSearchSelect(selectedLine int) error {
 	}
 	return nil
 }
+
+func (gui *Gui) handleRenameBranch(g *gocui.Gui, v *gocui.View) error {
+	branch := gui.getSelectedBranch()
+	if branch == nil {
+		return nil
+	}
+
+	promptForNewName := func() error {
+		return gui.createPromptPanel(g, v, gui.Tr.SLocalize("NewBranchNamePrompt")+" "+branch.Name+":", "", func(g *gocui.Gui, v *gocui.View) error {
+			newName := gui.trimmedContent(v)
+			if err := gui.GitCommand.RenameBranch(branch.Name, newName); err != nil {
+				return gui.createErrorPanel(gui.g, err.Error())
+			}
+			// need to checkout so that the branch shows up in our reflog and therefore
+			// doesn't get lost among all the other branches when we switch to something else
+			if err := gui.GitCommand.Checkout(newName, false); err != nil {
+				return gui.createErrorPanel(gui.g, err.Error())
+			}
+
+			return gui.refreshBranches(gui.g)
+		})
+	}
+
+	// I could do an explicit check here for whether the branch is tracking a remote branch
+	// but if we've selected it we'll already know that via Pullables and Pullables.
+	// Bit of a hack but I'm lazy.
+	notTrackingRemote := branch.Pullables == "?"
+	if notTrackingRemote {
+		return promptForNewName()
+	}
+	return gui.createConfirmationPanel(gui.g, v, true, gui.Tr.SLocalize("renameBranch"), gui.Tr.SLocalize("RenameBranchWarning"), func(_g *gocui.Gui, _v *gocui.View) error {
+		return promptForNewName()
+	}, nil)
+}
+
+func (gui *Gui) currentBranch() *commands.Branch {
+	return gui.State.Branches[0]
+}
diff --git a/pkg/gui/files_panel.go b/pkg/gui/files_panel.go
index 0b761bff7..bcfeccaa1 100644
--- a/pkg/gui/files_panel.go
+++ b/pkg/gui/files_panel.go
@@ -399,24 +399,20 @@ func (gui *Gui) catSelectedFile(g *gocui.Gui) (string, error) {
 
 func (gui *Gui) handlePullFiles(g *gocui.Gui, v *gocui.View) error {
 	// if we have no upstream branch we need to set that first
-	_, pullables := gui.GitCommand.GetCurrentBranchUpstreamDifferenceCount()
-	currentBranchName, err := gui.GitCommand.CurrentBranchName()
-	if err != nil {
-		return err
-	}
-	if pullables == "?" {
+	currentBranch := gui.currentBranch()
+	if currentBranch.Pullables == "?" {
 		// see if we have this branch in our config with an upstream
 		conf, err := gui.GitCommand.Repo.Config()
 		if err != nil {
 			return gui.createErrorPanel(gui.g, err.Error())
 		}
 		for branchName, branch := range conf.Branches {
-			if branchName == currentBranchName {
+			if branchName == currentBranch.Name {
 				return gui.pullFiles(v, fmt.Sprintf("%s %s", branch.Remote, branchName))
 			}
 		}
 
-		return gui.createPromptPanel(g, v, gui.Tr.SLocalize("EnterUpstream"), "origin/"+currentBranchName, func(g *gocui.Gui, v *gocui.View) error {
+		return gui.createPromptPanel(g, v, gui.Tr.SLocalize("EnterUpstream"), "origin/"+currentBranch.Name, func(g *gocui.Gui, v *gocui.View) error {
 			upstream := gui.trimmedContent(v)
 			if err := gui.GitCommand.SetUpstreamBranch(upstream); err != nil {
 				errorMessage := err.Error()
@@ -467,28 +463,24 @@ func (gui *Gui) pushWithForceFlag(g *gocui.Gui, v *gocui.View, force bool, upstr
 
 func (gui *Gui) pushFiles(g *gocui.Gui, v *gocui.View) error {
 	// if we have pullables we'll ask if the user wants to force push
-	_, pullables := gui.GitCommand.GetCurrentBranchUpstreamDifferenceCount()
-	currentBranchName, err := gui.GitCommand.CurrentBranchName()
-	if err != nil {
-		return err
-	}
+	currentBranch := gui.currentBranch()
 
-	if pullables == "?" {
+	if currentBranch.Pullables == "?" {
 		// see if we have this branch in our config with an upstream
 		conf, err := gui.GitCommand.Repo.Config()
 		if err != nil {
 			return gui.createErrorPanel(gui.g, err.Error())
 		}
 		for branchName, branch := range conf.Branches {
-			if branchName == currentBranchName {
+			if branchName == currentBranch.Name {
 				return gui.pushWithForceFlag(g, v, false, "", fmt.Sprintf("%s %s", branch.Remote, branchName))
 			}
 		}
 
-		return gui.createPromptPanel(g, v, gui.Tr.SLocalize("EnterUpstream"), "origin "+currentBranchName, func(g *gocui.Gui, v *gocui.View) error {
+		return gui.createPromptPanel(g, v, gui.Tr.SLocalize("EnterUpstream"), "origin "+currentBranch.Name, func(g *gocui.Gui, v *gocui.View) error {
 			return gui.pushWithForceFlag(g, v, false, gui.trimmedContent(v), "")
 		})
-	} else if pullables == "0" {
+	} else if currentBranch.Pullables == "0" {
 		return gui.pushWithForceFlag(g, v, false, "", "")
 	}
 	return gui.createConfirmationPanel(g, v, true, gui.Tr.SLocalize("ForcePush"), gui.Tr.SLocalize("ForcePushPrompt"), func(g *gocui.Gui, v *gocui.View) error {
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index 2204e3ac4..42f4bfc75 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -24,7 +24,6 @@ import (
 	"github.com/jesseduffield/gocui"
 	"github.com/jesseduffield/lazygit/pkg/commands"
 	"github.com/jesseduffield/lazygit/pkg/config"
-	"github.com/jesseduffield/lazygit/pkg/gui/presentation"
 	"github.com/jesseduffield/lazygit/pkg/i18n"
 	"github.com/jesseduffield/lazygit/pkg/tasks"
 	"github.com/jesseduffield/lazygit/pkg/theme"
@@ -278,6 +277,10 @@ func (gui *Gui) nextScreenMode(g *gocui.Gui, v *gocui.View) error {
 	if err := gui.refreshCommitsViewWithSelection(); err != nil {
 		return err
 	}
+	// same with branches
+	if err := gui.refreshBranchesViewWithSelection(); err != nil {
+		return err
+	}
 
 	return nil
 }
@@ -288,6 +291,10 @@ func (gui *Gui) prevScreenMode(g *gocui.Gui, v *gocui.View) error {
 	if err := gui.refreshCommitsViewWithSelection(); err != nil {
 		return err
 	}
+	// same with branches
+	if err := gui.refreshBranchesViewWithSelection(); err != nil {
+		return err
+	}
 
 	return nil
 }
@@ -386,12 +393,6 @@ func (gui *Gui) onFocusLost(v *gocui.View, newView *gocui.View) error {
 		}
 	}
 	switch v.Name() {
-	case "branches":
-		if v.Context == "local-branches" {
-			// This stops the branches panel from showing the upstream/downstream changes to the selected branch, when it loses focus
-			displayStrings := presentation.GetBranchListDisplayStrings(gui.State.Branches, false, -1)
-			gui.renderDisplayStrings(gui.getBranchesView(), displayStrings)
-		}
 	case "main":
 		// if we have lost focus to a first-class panel, we need to do some cleanup
 		gui.changeMainViewsContext("normal")
diff --git a/pkg/gui/keybindings.go b/pkg/gui/keybindings.go
index 0fb196aa3..258409e09 100644
--- a/pkg/gui/keybindings.go
+++ b/pkg/gui/keybindings.go
@@ -577,6 +577,14 @@ func (gui *Gui) GetInitialKeybindings() []*Binding {
 		},
 		{
 			ViewName:    "branches",
+			Contexts:    []string{"local-branches"},
+			Key:         gui.getKey("branches.renameBranch"),
+			Modifier:    gocui.ModNone,
+			Handler:     gui.handleRenameBranch,
+			Description: gui.Tr.SLocalize("viewResetOptions"),
+		},
+		{
+			ViewName:    "branches",
 			Contexts:    []string{"tags"},
 			Key:         gui.getKey("universal.select"),
 			Modifier:    gocui.ModNone,
diff --git a/pkg/gui/presentation/branches.go b/pkg/gui/presentation/branches.go
index 3b918116b..4283ee1dd 100644
--- a/pkg/gui/presentation/branches.go
+++ b/pkg/gui/presentation/branches.go
@@ -10,25 +10,38 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
-func GetBranchListDisplayStrings(branches []*commands.Branch, isFocused bool, selectedLine int) [][]string {
+func GetBranchListDisplayStrings(branches []*commands.Branch, fullDescription bool) [][]string {
 	lines := make([][]string, len(branches))
 
 	for i := range branches {
-		showUpstreamDifferences := isFocused && i == selectedLine
-		lines[i] = getBranchDisplayStrings(branches[i], showUpstreamDifferences)
+		lines[i] = getBranchDisplayStrings(branches[i], fullDescription)
 	}
 
 	return lines
 }
 
 // getBranchDisplayStrings returns the display string of branch
-func getBranchDisplayStrings(b *commands.Branch, showUpstreamDifferences bool) []string {
+func getBranchDisplayStrings(b *commands.Branch, fullDescription bool) []string {
 	displayName := utils.ColoredString(b.Name, GetBranchColor(b.Name))
-	if showUpstreamDifferences && b.Pushables != "" && b.Pullables != "" {
-		displayName = fmt.Sprintf("%s ↑%s↓%s", displayName, b.Pushables, b.Pullables)
+	if b.Pushables != "" && b.Pullables != "" && b.Pushables != "?" && b.Pullables != "?" {
+		trackColor := color.FgYellow
+		if b.Pushables == "0" && b.Pullables == "0" {
+			trackColor = color.FgGreen
+		}
+		track := utils.ColoredString(fmt.Sprintf("↑%s↓%s", b.Pushables, b.Pullables), trackColor)
+		displayName = fmt.Sprintf("%s %s", displayName, track)
 	}
 
-	return []string{b.Recency, displayName}
+	recencyColor := color.FgCyan
+	if b.Recency == "  *" {
+		recencyColor = color.FgGreen
+	}
+
+	if fullDescription {
+		return []string{utils.ColoredString(b.Recency, recencyColor), displayName, utils.ColoredString(b.UpstreamName, color.FgYellow)}
+	}
+
+	return []string{utils.ColoredString(b.Recency, recencyColor), displayName}
 }
 
 // GetBranchColor branch color
diff --git a/pkg/gui/reset_menu_panel.go b/pkg/gui/reset_menu_panel.go
index 5e9cc0068..1920aeb1e 100644
--- a/pkg/gui/reset_menu_panel.go
+++ b/pkg/gui/reset_menu_panel.go
@@ -36,6 +36,9 @@ func (gui *Gui) createResetMenu(ref string) error {
 				if err := gui.refreshFiles(); err != nil {
 					return err
 				}
+				if err := gui.refreshBranches(gui.g); err != nil {
+					return err
+				}
 				if err := gui.resetOrigin(gui.getCommitsView()); err != nil {
 					return err
 				}
diff --git a/pkg/gui/status_panel.go b/pkg/gui/status_panel.go
index 4c69e9b30..d7561d9d7 100644
--- a/pkg/gui/status_panel.go
+++ b/pkg/gui/status_panel.go
@@ -22,11 +22,20 @@ func (gui *Gui) refreshStatus(g *gocui.Gui) error {
 	// contents end up cleared
 	g.Update(func(*gocui.Gui) error {
 		v.Clear()
+		// TODO: base this off of the current branch
 		state.pushables, state.pullables = gui.GitCommand.GetCurrentBranchUpstreamDifferenceCount()
 		if err := gui.updateWorkTreeState(); err != nil {
 			return err
 		}
-		status := fmt.Sprintf("↑%s↓%s", state.pushables, state.pullables)
+
+		trackColor := color.FgYellow
+		if state.pushables == "0" && state.pullables == "0" {
+			trackColor = color.FgGreen
+		} else if state.pushables == "?" && state.pullables == "?" {
+			trackColor = color.FgRed
+		}
+
+		status := utils.ColoredString(fmt.Sprintf("↑%s↓%s", state.pushables, state.pullables), trackColor)
 		branches := gui.State.Branches
 
 		if gui.State.WorkingTreeState != "normal" {