support setting upstream

6 files changed, 55 insertions, 5 deletions

diff --git a/pkg/commands/git.go b/pkg/commands/git.go
index c2e12491d..9dda1b2d2 100644
--- a/pkg/commands/git.go
+++ b/pkg/commands/git.go
@@ -1087,3 +1087,7 @@ func (c *GitCommand) IsHeadDetached() bool {
 func (c *GitCommand) DeleteRemoteBranch(remoteName string, branchName string) error {
 	return c.OSCommand.RunCommand(fmt.Sprintf("git push %s --delete %s", remoteName, branchName))
 }
+
+func (c *GitCommand) SetBranchUpstream(remoteName string, remoteBranchName string, branchName string) error {
+	return c.OSCommand.RunCommand(fmt.Sprintf("git branch --set-upstream-to=%s/%s %s", remoteName, remoteBranchName, branchName))
+}
diff --git a/pkg/gui/branches_panel.go b/pkg/gui/branches_panel.go
index 28f16562a..970e383c4 100644
--- a/pkg/gui/branches_panel.go
+++ b/pkg/gui/branches_panel.go
@@ -212,8 +212,16 @@ func (gui *Gui) handleCheckoutByName(g *gocui.Gui, v *gocui.View) error {
 	return nil
 }
 
+func (gui *Gui) getCheckedOutBranch() *commands.Branch {
+	if len(gui.State.Branches) == 0 {
+		return nil
+	}
+
+	return gui.State.Branches[0]
+}
+
 func (gui *Gui) handleNewBranch(g *gocui.Gui, v *gocui.View) error {
-	branch := gui.State.Branches[0]
+	branch := gui.getCheckedOutBranch()
 	message := gui.Tr.TemplateLocalize(
 		"NewBranchNameBranchOff",
 		Teml{
@@ -243,7 +251,7 @@ func (gui *Gui) deleteBranch(g *gocui.Gui, v *gocui.View, force bool) error {
 	if selectedBranch == nil {
 		return nil
 	}
-	checkedOutBranch := gui.State.Branches[0]
+	checkedOutBranch := gui.getCheckedOutBranch()
 	if checkedOutBranch.Name == selectedBranch.Name {
 		return gui.createErrorPanel(g, gui.Tr.SLocalize("CantDeleteCheckOutBranch"))
 	}
@@ -280,7 +288,7 @@ func (gui *Gui) mergeBranchIntoCheckedOutBranch(branchName string) error {
 	if gui.GitCommand.IsHeadDetached() {
 		return gui.createErrorPanel(gui.g, "Cannot merge branch in detached head state. You might have checked out a commit directly or a remote branch, in which case you should checkout the local branch you want to be on")
 	}
-	checkedOutBranchName := gui.State.Branches[0].Name
+	checkedOutBranchName := gui.getCheckedOutBranch().Name
 	if checkedOutBranchName == branchName {
 		return gui.createErrorPanel(gui.g, gui.Tr.SLocalize("CantMergeBranchIntoItself"))
 	}
@@ -310,7 +318,7 @@ func (gui *Gui) handleRebaseOntoLocalBranch(g *gocui.Gui, v *gocui.View) error {
 }
 
 func (gui *Gui) handleRebaseOntoBranch(selectedBranchName string) error {
-	checkedOutBranch := gui.State.Branches[0].Name
+	checkedOutBranch := gui.getCheckedOutBranch().Name
 	if selectedBranchName == checkedOutBranch {
 		return gui.createErrorPanel(gui.g, gui.Tr.SLocalize("CantRebaseOntoSelf"))
 	}
diff --git a/pkg/gui/files_panel.go b/pkg/gui/files_panel.go
index 57df9def9..deffad322 100644
--- a/pkg/gui/files_panel.go
+++ b/pkg/gui/files_panel.go
@@ -451,7 +451,7 @@ func (gui *Gui) pushWithForceFlag(g *gocui.Gui, v *gocui.View, force bool, upstr
 	}
 	go func() {
 		unamePassOpend := false
-		branchName := gui.State.Branches[0].Name
+		branchName := gui.getCheckedOutBranch().Name
 		err := gui.GitCommand.Push(branchName, force, upstream, func(passOrUname string) string {
 			unamePassOpend = true
 			return gui.waitForPassUname(g, v, passOrUname)
diff --git a/pkg/gui/keybindings.go b/pkg/gui/keybindings.go
index 85678b6ad..85d7c5faf 100644
--- a/pkg/gui/keybindings.go
+++ b/pkg/gui/keybindings.go
@@ -1090,6 +1090,14 @@ func (gui *Gui) GetInitialKeybindings() []*Binding {
 			Description: gui.Tr.SLocalize("rebaseBranch"),
 		},
 		{
+			ViewName:    "branches",
+			Contexts:    []string{"remote-branches"},
+			Key:         'u',
+			Modifier:    gocui.ModNone,
+			Handler:     gui.handleSetBranchUpstream,
+			Description: gui.Tr.SLocalize("setUpstream"),
+		},
+		{
 			ViewName: "commits",
 			Key:      gocui.MouseLeft,
 			Modifier: gocui.ModNone,
diff --git a/pkg/gui/remote_branches_panel.go b/pkg/gui/remote_branches_panel.go
index c1534c89a..e4dfefa26 100644
--- a/pkg/gui/remote_branches_panel.go
+++ b/pkg/gui/remote_branches_panel.go
@@ -119,3 +119,24 @@ func (gui *Gui) handleRebaseOntoRemoteBranch(g *gocui.Gui, v *gocui.View) error
 	selectedBranchName := gui.getSelectedRemoteBranch().Name
 	return gui.handleRebaseOntoBranch(selectedBranchName)
 }
+
+func (gui *Gui) handleSetBranchUpstream(g *gocui.Gui, v *gocui.View) error {
+	selectedBranch := gui.getSelectedRemoteBranch()
+	checkedOutBranch := gui.getCheckedOutBranch()
+
+	message := gui.Tr.TemplateLocalize(
+		"SetUpstreamMessage",
+		Teml{
+			"checkedOut": checkedOutBranch.Name,
+			"selected":   selectedBranch.RemoteName + "/" + selectedBranch.Name,
+		},
+	)
+
+	return gui.createConfirmationPanel(g, v, true, gui.Tr.SLocalize("SetUpstreamTitle"), message, func(*gocui.Gui, *gocui.View) error {
+		if err := gui.GitCommand.SetBranchUpstream(selectedBranch.RemoteName, selectedBranch.Name, checkedOutBranch.Name); err != nil {
+			return err
+		}
+
+		return gui.refreshSidePanels(gui.g)
+	}, nil)
+}
diff --git a/pkg/i18n/english.go b/pkg/i18n/english.go
index 736bf9162..d39b4b279 100644
--- a/pkg/i18n/english.go
+++ b/pkg/i18n/english.go
@@ -870,6 +870,15 @@ func addEnglish(i18nObject *i18n.Bundle) error {
 		}, &i18n.Message{
 			ID:    "DeleteRemoteBranchMessage",
 			Other: "Are you sure you want to delete remote branch",
+		}, &i18n.Message{
+			ID:    "setUpstream",
+			Other: "set as upstream of checked-out branch",
+		}, &i18n.Message{
+			ID:    "SetUpstreamTitle",
+			Other: "Set upstream branch",
+		}, &i18n.Message{
+			ID:    "SetUpstreamMessage",
+			Other: "Are you sure you want to set the upstream branch of '{{.checkedOut}}' to '{{.selected}}'",
 		},
 	)
 }