Offer autostash option when creating new branch (#3871)

- **PR Description**

Resolves https://github.com/jesseduffield/lazygit/issues/3866

My attempt at creating a shared abstraction felt like it was more
indirection than it was worth so I ended up going with the code
duplication approach.

pkg/gui/controllers/helpers/refs_helper.go

@@ -44,15 +44,15 @@ func (self *RefsHelper) CheckoutRef(ref string, options types.CheckoutRefOptions
 
 	cmdOptions := git_commands.CheckoutOptions{Force: false, EnvVars: options.EnvVars}
 
-	onSuccess := func() {
+	refresh := func() {
 		self.c.Contexts().Branches.SetSelection(0)
 		self.c.Contexts().ReflogCommits.SetSelection(0)
 		self.c.Contexts().LocalCommits.SetSelection(0)
 		// loading a heap of commits is slow so we limit them whenever doing a reset
 		self.c.Contexts().LocalCommits.SetLimitCommits(true)
-	}
 
-	refreshOptions := types.RefreshOptions{Mode: types.BLOCK_UI, KeepBranchSelectionIndex: true}
+		_ = self.c.Refresh(types.RefreshOptions{Mode: types.BLOCK_UI, KeepBranchSelectionIndex: true})
+	}
 
 	localBranch, found := lo.Find(self.c.Model().Branches, func(branch *models.Branch) bool {
 		return branch.Name == ref
@@ -74,7 +74,7 @@ func (self *RefsHelper) CheckoutRef(ref string, options types.CheckoutRefOptions
 				return options.OnRefNotFound(ref)
 			}
 
-			if strings.Contains(err.Error(), "Please commit your changes or stash them before you switch branch") {
+			if IsSwitchBranchUncommitedChangesError(err) {
 				// offer to autostash changes
 				self.c.OnUIThread(func() error {
 					// (Before showing the prompt, render again to remove the inline status)
@@ -90,15 +90,10 @@ func (self *RefsHelper) CheckoutRef(ref string, options types.CheckoutRefOptions
 								if err := self.c.Git().Branch.Checkout(ref, cmdOptions); err != nil {
 									return err
 								}
-
-								onSuccess()
-								if err := self.c.Git().Stash.Pop(0); err != nil {
-									if err := self.c.Refresh(refreshOptions); err != nil {
-										return err
-									}
-									return err
-								}
-								return self.c.Refresh(refreshOptions)
+								err := self.c.Git().Stash.Pop(0)
+								// Branch switch successful so re-render the UI even if the pop operation failed (e.g. conflict).
+								refresh()
+								return err
 							})
 						},
 					})
@@ -108,9 +103,9 @@ func (self *RefsHelper) CheckoutRef(ref string, options types.CheckoutRefOptions
 
 			return err
 		}
-		onSuccess()
 
-		return self.c.Refresh(refreshOptions)
+		refresh()
+		return nil
 	})
 }
 
@@ -288,6 +283,19 @@ func (self *RefsHelper) NewBranch(from string, fromFormattedName string, suggest
 		suggestedBranchName = self.c.UserConfig().Git.BranchPrefix
 	}
 
+	refresh := func() error {
+		if self.c.Context().Current() != self.c.Contexts().Branches {
+			if err := self.c.Context().Push(self.c.Contexts().Branches); err != nil {
+				return err
+			}
+		}
+
+		self.c.Contexts().LocalCommits.SetSelection(0)
+		self.c.Contexts().Branches.SetSelection(0)
+
+		return self.c.Refresh(types.RefreshOptions{Mode: types.BLOCK_UI, KeepBranchSelectionIndex: true})
+	}
+
 	return self.c.Prompt(types.PromptOpts{
 		Title:          message,
 		InitialContent: suggestedBranchName,
@@ -299,19 +307,34 @@ func (self *RefsHelper) NewBranch(from string, fromFormattedName string, suggest
 				newBranchFunc = self.c.Git().Branch.NewWithoutTracking
 			}
 			if err := newBranchFunc(newBranchName, from); err != nil {
+				if IsSwitchBranchUncommitedChangesError(err) {
+					// offer to autostash changes
+					return self.c.Confirm(types.ConfirmOpts{
+						Title:  self.c.Tr.AutoStashTitle,
+						Prompt: self.c.Tr.AutoStashPrompt,
+						HandleConfirm: func() error {
+							if err := self.c.Git().Stash.Push(self.c.Tr.StashPrefix + newBranchName); err != nil {
+								return err
+							}
+							if err := newBranchFunc(newBranchName, from); err != nil {
+								return err
+							}
+							popErr := self.c.Git().Stash.Pop(0)
+							// Branch switch successful so re-render the UI even if the pop operation failed (e.g. conflict).
+							refreshError := refresh()
+							if popErr != nil {
+								// An error from pop is the more important one to report to the user
+								return popErr
+							}
+							return refreshError
+						},
+					})
+				}
+
 				return err
 			}
 
-			if self.c.Context().Current() != self.c.Contexts().Branches {
-				if err := self.c.Context().Push(self.c.Contexts().Branches); err != nil {
-					return err
-				}
-			}
-
-			self.c.Contexts().LocalCommits.SetSelection(0)
-			self.c.Contexts().Branches.SetSelection(0)
-
-			return self.c.Refresh(types.RefreshOptions{Mode: types.BLOCK_UI, KeepBranchSelectionIndex: true})
+			return refresh()
 		},
 	})
 }
@@ -339,3 +362,7 @@ func (self *RefsHelper) ParseRemoteBranchName(fullBranchName string) (string, st
 
 	return remoteName, branchName, true
 }
+
+func IsSwitchBranchUncommitedChangesError(err error) bool {
+	return strings.Contains(err.Error(), "Please commit your changes or stash them before you switch branch")
+}

pkg/integration/tests/branch/checkout_autostash.go

@@ -0,0 +1,54 @@
+package branch
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var CheckoutAutostash = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Check out a branch that requires performing autostash",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig:  func(config *config.AppConfig) {},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file", "a\n\nb")
+		shell.Commit("add file")
+		shell.UpdateFileAndAdd("file", "a\n\nc")
+		shell.Commit("edit last line")
+
+		shell.Checkout("HEAD^")
+		shell.UpdateFile("file", "b\n\nb")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			Lines(
+				Contains("file"),
+			)
+
+		t.Views().Branches().
+			Focus().
+			Lines(
+				MatchesRegexp(`\*.*HEAD`).IsSelected(),
+				Contains("master"),
+			).
+			NavigateToLine(Contains("master")).
+			PressPrimaryAction()
+
+		t.ExpectPopup().Confirmation().
+			Title(Contains("Autostash?")).
+			Content(Contains("You must stash and pop your changes to bring them across. Do this automatically? (enter/esc)")).
+			Confirm()
+
+		t.Views().Branches().
+			Lines(
+				Contains("master").IsSelected(),
+			)
+
+		t.Git().CurrentBranchName("master")
+
+		t.Views().Files().
+			Lines(
+				Contains("file"),
+			)
+	},
+})

pkg/integration/tests/branch/new_branch_autostash.go

@@ -0,0 +1,60 @@
+package branch
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var NewBranchAutostash = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Create a new branch that requires performing autostash",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig:  func(config *config.AppConfig) {},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file", "a\n\nb")
+		shell.Commit("add file")
+		shell.UpdateFileAndAdd("file", "a\n\nc")
+		shell.Commit("edit last line")
+
+		shell.Checkout("HEAD^")
+		shell.UpdateFile("file", "b\n\nb")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			Lines(
+				Contains("file"),
+			)
+
+		t.Views().Branches().
+			Focus().
+			Lines(
+				MatchesRegexp(`\*.*HEAD`).IsSelected(),
+				Contains("master"),
+			).
+			NavigateToLine(Contains("master")).
+			Press(keys.Universal.New)
+
+		t.ExpectPopup().Prompt().
+			Title(Contains("New branch name (branch is off of 'master')")).
+			Type("new-branch").
+			Confirm()
+
+		t.ExpectPopup().Confirmation().
+			Title(Contains("Autostash?")).
+			Content(Contains("You must stash and pop your changes to bring them across. Do this automatically? (enter/esc)")).
+			Confirm()
+
+		t.Views().Branches().
+			Lines(
+				Contains("new-branch").IsSelected(),
+				Contains("master"),
+			)
+
+		t.Git().CurrentBranchName("new-branch")
+
+		t.Views().Files().
+			Lines(
+				Contains("file"),
+			)
+	},
+})

pkg/integration/tests/test_list.go

@@ -38,11 +38,13 @@ var tests = []*components.IntegrationTest{
 	bisect.ChooseTerms,
 	bisect.FromOtherBranch,
 	bisect.Skip,
+	branch.CheckoutAutostash,
 	branch.CheckoutByName,
 	branch.CreateTag,
 	branch.Delete,
 	branch.DeleteRemoteBranchWithCredentialPrompt,
 	branch.DetachedHead,
+	branch.NewBranchAutostash,
 	branch.NewBranchFromRemoteTrackingDifferentName,
 	branch.NewBranchFromRemoteTrackingSameName,
 	branch.OpenPullRequestNoUpstream,