support discarding changes in dir

4 files changed, 106 insertions, 26 deletions

diff --git a/pkg/commands/files.go b/pkg/commands/files.go
index c73a898b4..740928759 100644
--- a/pkg/commands/files.go
+++ b/pkg/commands/files.go
@@ -2,6 +2,7 @@ package commands
 
 import (
 	"fmt"
+	"os"
 	"os/exec"
 	"path/filepath"
 	"strings"
@@ -140,6 +141,47 @@ func (c *GitCommand) DiscardAllFileChanges(file *models.File) error {
 	return c.DiscardUnstagedFileChanges(file)
 }
 
+func (c *GitCommand) DiscardAllDirChanges(node *models.StatusLineNode) error {
+	if err := c.RemoveUntrackedDirFiles(node); err != nil {
+		return err
+	}
+
+	quotedPath := c.OSCommand.Quote(node.GetPath())
+	if err := c.OSCommand.RunCommand("git checkout HEAD -- %s", quotedPath); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (c *GitCommand) DiscardUnstagedDirChanges(node *models.StatusLineNode) error {
+	if err := c.RemoveUntrackedDirFiles(node); err != nil {
+		return err
+	}
+
+	quotedPath := c.OSCommand.Quote(node.GetPath())
+	if err := c.OSCommand.RunCommand("git checkout -- %s", quotedPath); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (c *GitCommand) RemoveUntrackedDirFiles(node *models.StatusLineNode) error {
+	untrackedFilePaths := node.GetPathsMatching(
+		func(n *models.StatusLineNode) bool { return n.File != nil && !n.File.GetIsTracked() },
+	)
+
+	for _, path := range untrackedFilePaths {
+		err := os.Remove(path)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // DiscardUnstagedFileChanges directly
 func (c *GitCommand) DiscardUnstagedFileChanges(file *models.File) error {
 	quotedFileName := c.OSCommand.Quote(file.Name)
diff --git a/pkg/commands/models/status_line_node.go b/pkg/commands/models/status_line_node.go
index 3ed1a573a..1a2471935 100644
--- a/pkg/commands/models/status_line_node.go
+++ b/pkg/commands/models/status_line_node.go
@@ -188,3 +188,18 @@ func (s *StatusLineNode) compressAux() *StatusLineNode {
 func (s *StatusLineNode) HasExactlyOneChild() bool {
 	return len(s.Children) == 1
 }
+
+// This ignores the root
+func (s *StatusLineNode) GetPathsMatching(test func(*StatusLineNode) bool) []string {
+	paths := []string{}
+
+	if test(s) {
+		paths = append(paths, s.GetPath())
+	}
+
+	for _, child := range s.Children {
+		paths = append(paths, child.GetPathsMatching(test)...)
+	}
+
+	return paths
+}
diff --git a/pkg/gui/discard_changes_menu_panel.go b/pkg/gui/discard_changes_menu_panel.go
index f81297479..da4e91e18 100644
--- a/pkg/gui/discard_changes_menu_panel.go
+++ b/pkg/gui/discard_changes_menu_panel.go
@@ -1,35 +1,18 @@
 package gui
 
-import (
-	"github.com/jesseduffield/gocui"
-)
-
-func (gui *Gui) handleCreateDiscardMenu(g *gocui.Gui, v *gocui.View) error {
-	file := gui.getSelectedFile()
-	if file == nil {
+func (gui *Gui) handleCreateDiscardMenu() error {
+	node := gui.getSelectedStatusNode()
+	if node == nil {
 		return nil
 	}
 
 	var menuItems []*menuItem
-
-	submodules := gui.State.Submodules
-	if file.IsSubmodule(submodules) {
-		submodule := file.SubmoduleConfig(submodules)
-
-		menuItems = []*menuItem{
-			{
-				displayString: gui.Tr.LcSubmoduleStashAndReset,
-				onPress: func() error {
-					return gui.handleResetSubmodule(submodule)
-				},
-			},
-		}
-	} else {
+	if node.File == nil {
 		menuItems = []*menuItem{
 			{
 				displayString: gui.Tr.LcDiscardAllChanges,
 				onPress: func() error {
-					if err := gui.GitCommand.DiscardAllFileChanges(file); err != nil {
+					if err := gui.GitCommand.DiscardAllDirChanges(node); err != nil {
 						return gui.surfaceError(err)
 					}
 					return gui.refreshSidePanels(refreshOptions{mode: ASYNC, scope: []int{FILES}})
@@ -37,11 +20,11 @@ func (gui *Gui) handleCreateDiscardMenu(g *gocui.Gui, v *gocui.View) error {
 			},
 		}
 
-		if file.HasStagedChanges && file.HasUnstagedChanges {
+		if node.GetHasStagedChanges() && node.GetHasUnstagedChanges() {
 			menuItems = append(menuItems, &menuItem{
 				displayString: gui.Tr.LcDiscardUnstagedChanges,
 				onPress: func() error {
-					if err := gui.GitCommand.DiscardUnstagedFileChanges(file); err != nil {
+					if err := gui.GitCommand.DiscardUnstagedDirChanges(node); err != nil {
 						return gui.surfaceError(err)
 					}
 
@@ -49,8 +32,48 @@ func (gui *Gui) handleCreateDiscardMenu(g *gocui.Gui, v *gocui.View) error {
 				},
 			})
 		}
+	} else {
+		file := node.File
+
+		submodules := gui.State.Submodules
+		if file.IsSubmodule(submodules) {
+			submodule := file.SubmoduleConfig(submodules)
+
+			menuItems = []*menuItem{
+				{
+					displayString: gui.Tr.LcSubmoduleStashAndReset,
+					onPress: func() error {
+						return gui.handleResetSubmodule(submodule)
+					},
+				},
+			}
+		} else {
+			menuItems = []*menuItem{
+				{
+					displayString: gui.Tr.LcDiscardAllChanges,
+					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.LcDiscardUnstagedChanges,
+					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.createMenu(file.Name, menuItems, createMenuOptions{showCancel: true})
+	return gui.createMenu(node.GetPath(), menuItems, createMenuOptions{showCancel: true})
 }
diff --git a/pkg/gui/keybindings.go b/pkg/gui/keybindings.go
index c01aad9ac..ea6e13e99 100644
--- a/pkg/gui/keybindings.go
+++ b/pkg/gui/keybindings.go
@@ -408,7 +408,7 @@ func (gui *Gui) GetInitialKeybindings() []*Binding {
 			ViewName:    "files",
 			Contexts:    []string{FILES_CONTEXT_KEY},
 			Key:         gui.getKey(config.Universal.Remove),
-			Handler:     gui.handleCreateDiscardMenu,
+			Handler:     gui.wrappedHandler(gui.handleCreateDiscardMenu),
 			Description: gui.Tr.LcViewDiscardOptions,
 			OpensMenu:   true,
 		},