support discarding changes in dir

1 files changed, 42 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)