allow explicitly managing focus

3 files changed, 58 insertions, 36 deletions

diff --git a/pkg/gui/commit_files_panel.go b/pkg/gui/commit_files_panel.go
index 505beda0e..15d073785 100644
--- a/pkg/gui/commit_files_panel.go
+++ b/pkg/gui/commit_files_panel.go
@@ -185,8 +185,9 @@ func (gui *Gui) enterCommitFile(selectedLineIdx int) error {
 
 	if gui.GitCommand.PatchManager.Active() && gui.GitCommand.PatchManager.To != commitFile.Parent {
 		return gui.ask(askOpts{
-			title:  gui.Tr.SLocalize("DiscardPatch"),
-			prompt: gui.Tr.SLocalize("DiscardPatchConfirm"),
+			title:               gui.Tr.SLocalize("DiscardPatch"),
+			prompt:              gui.Tr.SLocalize("DiscardPatchConfirm"),
+			handlersManageFocus: true,
 			handleConfirm: func() error {
 				gui.GitCommand.PatchManager.Reset()
 				return enterTheFile(selectedLineIdx)
diff --git a/pkg/gui/confirmation_panel.go b/pkg/gui/confirmation_panel.go
index e9103b42a..ed8a0137f 100644
--- a/pkg/gui/confirmation_panel.go
+++ b/pkg/gui/confirmation_panel.go
@@ -17,21 +17,23 @@ import (
 
 type createPopupPanelOpts struct {
 	hasLoader           bool
-	returnFocusOnClose  bool
 	editable            bool
 	title               string
 	prompt              string
 	handleConfirm       func() error
 	handleConfirmPrompt func(string) error
 	handleClose         func() error
+
+	// when handlersManageFocus is true, do not return from the confirmation context automatically. It's expected that the handlers will manage focus, whether that means switching to another context, or manually returning the context.
+	handlersManageFocus bool
 }
 
 type askOpts struct {
-	returnFocusOnClose bool
-	title              string
-	prompt             string
-	handleConfirm      func() error
-	handleClose        func() error
+	title               string
+	prompt              string
+	handleConfirm       func() error
+	handleClose         func() error
+	handlersManageFocus bool
 }
 
 func (gui *Gui) createLoaderPanel(currentView *gocui.View, prompt string) error {
@@ -43,11 +45,11 @@ func (gui *Gui) createLoaderPanel(currentView *gocui.View, prompt string) error
 
 func (gui *Gui) ask(opts askOpts) error {
 	return gui.createPopupPanel(createPopupPanelOpts{
-		title:  opts.title,
-		prompt: opts.prompt,
-
-		handleConfirm: opts.handleConfirm,
-		handleClose:   opts.handleClose,
+		title:               opts.title,
+		prompt:              opts.prompt,
+		handleConfirm:       opts.handleConfirm,
+		handleClose:         opts.handleClose,
+		handlersManageFocus: opts.handlersManageFocus,
 	})
 }
 
@@ -60,47 +62,60 @@ func (gui *Gui) prompt(title string, initialContent string, handleConfirm func(s
 	})
 }
 
-func (gui *Gui) wrappedConfirmationFunction(function func() error, returnFocusOnClose bool) func(*gocui.Gui, *gocui.View) error {
+func (gui *Gui) wrappedConfirmationFunction(handlersManageFocus bool, function func() error) func(*gocui.Gui, *gocui.View) error {
 	return func(g *gocui.Gui, v *gocui.View) error {
-
 		if function != nil {
 			if err := function(); err != nil {
 				return err
 			}
 		}
 
-		return gui.closeConfirmationPrompt(returnFocusOnClose)
+		if err := gui.closeConfirmationPrompt(handlersManageFocus); err != nil {
+			return err
+		}
+
+		return nil
 	}
 }
 
-func (gui *Gui) wrappedPromptConfirmationFunction(function func(string) error, returnFocusOnClose bool) func(*gocui.Gui, *gocui.View) error {
+func (gui *Gui) wrappedPromptConfirmationFunction(handlersManageFocus bool, function func(string) error) func(*gocui.Gui, *gocui.View) error {
 	return func(g *gocui.Gui, v *gocui.View) error {
-
 		if function != nil {
 			if err := function(v.Buffer()); err != nil {
 				return gui.surfaceError(err)
 			}
 		}
 
-		return gui.closeConfirmationPrompt(returnFocusOnClose)
+		if err := gui.closeConfirmationPrompt(handlersManageFocus); err != nil {
+			return err
+		}
+
+		return nil
 	}
 }
 
-func (gui *Gui) closeConfirmationPrompt(returnFocusOnClose bool) error {
+func (gui *Gui) deleteConfirmationView() {
+	gui.g.DeleteKeybinding("confirmation", gui.getKey("universal.confirm"), gocui.ModNone)
+	gui.g.DeleteKeybinding("confirmation", gui.getKey("universal.confirm-alt1"), gocui.ModNone)
+	gui.g.DeleteKeybinding("confirmation", gui.getKey("universal.return"), gocui.ModNone)
+
+	_ = gui.g.DeleteView("confirmation")
+}
+
+func (gui *Gui) closeConfirmationPrompt(handlersManageFocus bool) error {
 	view := gui.getConfirmationView()
 	if view == nil {
 		return nil // if it's already been closed we can just return
 	}
-	view.Editable = false
-	if err := gui.returnFromContext(); err != nil {
-		return err
-	}
 
-	gui.g.DeleteKeybinding("confirmation", gui.getKey("universal.confirm"), gocui.ModNone)
-	gui.g.DeleteKeybinding("confirmation", gui.getKey("universal.confirm-alt1"), gocui.ModNone)
-	gui.g.DeleteKeybinding("confirmation", gui.getKey("universal.return"), gocui.ModNone)
+	if !handlersManageFocus {
+		if err := gui.returnFromContext(); err != nil {
+			return err
+		}
+	}
 
-	return gui.g.DeleteView("confirmation")
+	gui.deleteConfirmationView()
+	return nil
 }
 
 func (gui *Gui) getMessageHeight(wrap bool, message string, width int) int {
@@ -164,9 +179,7 @@ func (gui *Gui) createPopupPanel(opts createPopupPanelOpts) error {
 	gui.g.Update(func(g *gocui.Gui) error {
 		// delete the existing confirmation panel if it exists
 		if view, _ := g.View("confirmation"); view != nil {
-			if err := gui.closeConfirmationPrompt(true); err != nil {
-				gui.Log.Error(err)
-			}
+			gui.deleteConfirmationView()
 		}
 		confirmationView, err := gui.prepareConfirmationPanel(opts.title, opts.prompt, opts.hasLoader)
 		if err != nil {
@@ -202,9 +215,9 @@ func (gui *Gui) setKeyBindings(opts createPopupPanelOpts) error {
 	gui.renderString("options", actions)
 	var onConfirm func(*gocui.Gui, *gocui.View) error
 	if opts.handleConfirmPrompt != nil {
-		onConfirm = gui.wrappedPromptConfirmationFunction(opts.handleConfirmPrompt, opts.returnFocusOnClose)
+		onConfirm = gui.wrappedPromptConfirmationFunction(opts.handlersManageFocus, opts.handleConfirmPrompt)
 	} else {
-		onConfirm = gui.wrappedConfirmationFunction(opts.handleConfirm, opts.returnFocusOnClose)
+		onConfirm = gui.wrappedConfirmationFunction(opts.handlersManageFocus, opts.handleConfirm)
 	}
 
 	if err := gui.g.SetKeybinding("confirmation", nil, gui.getKey("universal.confirm"), gocui.ModNone, onConfirm); err != nil {
@@ -214,7 +227,7 @@ func (gui *Gui) setKeyBindings(opts createPopupPanelOpts) error {
 		return err
 	}
 
-	return gui.g.SetKeybinding("confirmation", nil, gui.getKey("universal.return"), gocui.ModNone, gui.wrappedConfirmationFunction(opts.handleClose, opts.returnFocusOnClose))
+	return gui.g.SetKeybinding("confirmation", nil, gui.getKey("universal.return"), gocui.ModNone, gui.wrappedConfirmationFunction(opts.handlersManageFocus, opts.handleClose))
 }
 
 func (gui *Gui) createErrorPanel(message string) error {
diff --git a/pkg/gui/staging_panel.go b/pkg/gui/staging_panel.go
index a9b351ba8..90e776024 100644
--- a/pkg/gui/staging_panel.go
+++ b/pkg/gui/staging_panel.go
@@ -105,11 +105,19 @@ func (gui *Gui) handleResetSelection(g *gocui.Gui, v *gocui.View) error {
 
 	if !gui.Config.GetUserConfig().GetBool("gui.skipUnstageLineWarning") {
 		return gui.ask(askOpts{
-			title:  gui.Tr.SLocalize("UnstageLinesTitle"),
-			prompt: gui.Tr.SLocalize("UnstageLinesPrompt"),
+			title:               gui.Tr.SLocalize("UnstageLinesTitle"),
+			prompt:              gui.Tr.SLocalize("UnstageLinesPrompt"),
+			handlersManageFocus: true,
 			handleConfirm: func() error {
+				if err := gui.switchContext(gui.Contexts.Staging.Context); err != nil {
+					return err
+				}
+
 				return gui.applySelection(true)
 			},
+			handleClose: func() error {
+				return gui.switchContext(gui.Contexts.Staging.Context)
+			},
 		})
 	} else {
 		return gui.applySelection(true)