support discarding changes in dir

2 files changed, 57 insertions, 0 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
+}