support discarding submodule changes

5 files changed, 81 insertions, 34 deletions

diff --git a/pkg/gui/discard_changes_menu_panel.go b/pkg/gui/discard_changes_menu_panel.go
index f8a0f1df2..2c4e8f759 100644
--- a/pkg/gui/discard_changes_menu_panel.go
+++ b/pkg/gui/discard_changes_menu_panel.go
@@ -2,42 +2,74 @@ package gui
 
 import (
 	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/commands"
 )
 
+func (gui *Gui) submoduleFromFile(file *commands.File) *commands.SubmoduleConfig {
+	for _, config := range gui.State.SubmoduleConfigs {
+		if config.Name == file.Name {
+			return config
+		}
+	}
+
+	return nil
+}
+
 func (gui *Gui) handleCreateDiscardMenu(g *gocui.Gui, v *gocui.View) error {
 	file := gui.getSelectedFile()
 	if file == nil {
 		return nil
 	}
 
-	if file.IsSubmodule {
-		// git submodule foreach '[[ "$name" == "renderers/chartify" ]] && git stash --include-untracked'
-		// git submodule update --force renderers/chartify
-	}
+	var menuItems []*menuItem
 
-	menuItems := []*menuItem{
-		{
-			displayString: gui.Tr.SLocalize("discardAllChanges"),
-			onPress: func() error {
-				if err := gui.GitCommand.DiscardAllFileChanges(file); err != nil {
-					return gui.surfaceError(err)
-				}
-				return gui.refreshSidePanels(refreshOptions{mode: ASYNC, scope: []int{FILES}})
+	submoduleConfigs := gui.State.SubmoduleConfigs
+	if file.IsSubmodule(submoduleConfigs) {
+		submoduleConfig := file.SubmoduleConfig(submoduleConfigs)
+
+		menuItems = []*menuItem{
+			{
+				displayString: gui.Tr.SLocalize("submoduleStashAndReset"),
+				onPress: func() error {
+					if err := gui.GitCommand.UnStageFile(file.Name, file.Tracked); err != nil {
+						return gui.surfaceError(err)
+					}
+					if err := gui.GitCommand.SubmoduleStash(submoduleConfig); err != nil {
+						return gui.surfaceError(err)
+					}
+					if err := gui.GitCommand.SubmoduleReset(submoduleConfig); err != nil {
+						return gui.surfaceError(err)
+					}
+					return gui.refreshSidePanels(refreshOptions{mode: ASYNC, scope: []int{FILES}})
+				},
 			},
-		},
-	}
+		}
+	} else {
+		menuItems = []*menuItem{
+			{
+				displayString: gui.Tr.SLocalize("discardAllChanges"),
+				onPress: func() error {
+					if err := gui.GitCommand.DiscardAllFileChanges(file); err != nil {
+						return gui.surfaceError(err)
+					}
+					return gui.refreshSidePanels(refreshOptions{mode: ASYNC, scope: []int{FILES}})
+				},
+			},
+		}
 
-	if file.HasStagedChanges && file.HasUnstagedChanges {
-		menuItems = append(menuItems, &menuItem{
-			displayString: gui.Tr.SLocalize("discardUnstagedChanges"),
-			onPress: func() error {
-				if err := gui.GitCommand.DiscardUnstagedFileChanges(file); err != nil {
-					return gui.surfaceError(err)
-				}
+		if file.HasStagedChanges && file.HasUnstagedChanges {
+			menuItems = append(menuItems, &menuItem{
+				displayString: gui.Tr.SLocalize("discardUnstagedChanges"),
+				onPress: func() error {
+					if err := gui.GitCommand.DiscardUnstagedFileChanges(file); err != nil {
+						return gui.surfaceError(err)
+					}
+
+					return gui.refreshSidePanels(refreshOptions{mode: ASYNC, scope: []int{FILES}})
+				},
+			})
+		}
 
-				return gui.refreshSidePanels(refreshOptions{mode: ASYNC, scope: []int{FILES}})
-			},
-		})
 	}
 
 	return gui.createMenu(file.Name, menuItems, createMenuOptions{showCancel: true})
diff --git a/pkg/gui/files_panel.go b/pkg/gui/files_panel.go
index aa01b4b36..c57f94082 100644
--- a/pkg/gui/files_panel.go
+++ b/pkg/gui/files_panel.go
@@ -93,6 +93,9 @@ func (gui *Gui) refreshFiles() error {
 		// if the filesView hasn't been instantiated yet we just return
 		return nil
 	}
+	if err := gui.refreshStateSubmoduleConfigs(); err != nil {
+		return err
+	}
 	if err := gui.refreshStateFiles(); err != nil {
 		return err
 	}
@@ -426,6 +429,17 @@ func (gui *Gui) refreshStateFiles() error {
 	return nil
 }
 
+func (gui *Gui) refreshStateSubmoduleConfigs() error {
+	configs, err := gui.GitCommand.GetSubmoduleConfigs()
+	if err != nil {
+		return err
+	}
+
+	gui.State.SubmoduleConfigs = configs
+
+	return nil
+}
+
 func (gui *Gui) handlePullFiles(g *gocui.Gui, v *gocui.View) error {
 	if gui.popupPanelFocused() {
 		return nil
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index db6624ec0..94634113b 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -271,11 +271,12 @@ type Modes struct {
 }
 
 type guiState struct {
-	Files        []*commands.File
-	Branches     []*commands.Branch
-	Commits      []*commands.Commit
-	StashEntries []*commands.StashEntry
-	CommitFiles  []*commands.CommitFile
+	Files            []*commands.File
+	SubmoduleConfigs []*commands.SubmoduleConfig
+	Branches         []*commands.Branch
+	Commits          []*commands.Commit
+	StashEntries     []*commands.StashEntry
+	CommitFiles      []*commands.CommitFile
 	// FilteredReflogCommits are the ones that appear in the reflog panel.
 	// when in filtering mode we only include the ones that match the given path
 	FilteredReflogCommits []*commands.Commit
diff --git a/pkg/gui/list_context.go b/pkg/gui/list_context.go
index 1ba7da80a..e2dec2a96 100644
--- a/pkg/gui/list_context.go
+++ b/pkg/gui/list_context.go
@@ -270,7 +270,7 @@ func (gui *Gui) filesListContext() *ListContext {
 		ResetMainViewOriginOnFocus: false,
 		Kind:                       SIDE_CONTEXT,
 		GetDisplayStrings: func() [][]string {
-			return presentation.GetFileListDisplayStrings(gui.State.Files, gui.State.Modes.Diffing.Ref)
+			return presentation.GetFileListDisplayStrings(gui.State.Files, gui.State.Modes.Diffing.Ref, gui.State.SubmoduleConfigs)
 		},
 		SelectedItem: func() (ListItem, bool) {
 			item := gui.getSelectedFile()
diff --git a/pkg/gui/presentation/files.go b/pkg/gui/presentation/files.go
index 6493e1024..f02aee91f 100644
--- a/pkg/gui/presentation/files.go
+++ b/pkg/gui/presentation/files.go
@@ -7,19 +7,19 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
-func GetFileListDisplayStrings(files []*commands.File, diffName string) [][]string {
+func GetFileListDisplayStrings(files []*commands.File, diffName string, submoduleConfigs []*commands.SubmoduleConfig) [][]string {
 	lines := make([][]string, len(files))
 
 	for i := range files {
 		diffed := files[i].Name == diffName
-		lines[i] = getFileDisplayStrings(files[i], diffed)
+		lines[i] = getFileDisplayStrings(files[i], diffed, submoduleConfigs)
 	}
 
 	return lines
 }
 
 // getFileDisplayStrings returns the display string of branch
-func getFileDisplayStrings(f *commands.File, diffed bool) []string {
+func getFileDisplayStrings(f *commands.File, diffed bool, submoduleConfigs []*commands.SubmoduleConfig) []string {
 	// potentially inefficient to be instantiating these color
 	// objects with each render
 	red := color.New(color.FgRed)
@@ -55,7 +55,7 @@ func getFileDisplayStrings(f *commands.File, diffed bool) []string {
 	output += secondCharCl.Sprint(secondChar)
 	output += restColor.Sprintf(" %s", f.Name)
 
-	if f.IsSubmodule {
+	if f.IsSubmodule(submoduleConfigs) {
 		output += utils.ColoredString(" (submodule)", theme.DefaultTextColor)
 	}