more file tree improvements

4 files changed, 28 insertions, 46 deletions

diff --git a/pkg/commands/files.go b/pkg/commands/files.go
index e7332dbbe..8244408cc 100644
--- a/pkg/commands/files.go
+++ b/pkg/commands/files.go
@@ -142,16 +142,8 @@ func (c *GitCommand) DiscardAllFileChanges(file *models.File) error {
 }
 
 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
+	// this could be more efficient but we would need to handle all the edge cases
+	return node.ForEachFile(c.DiscardAllFileChanges)
 }
 
 func (c *GitCommand) DiscardUnstagedDirChanges(node *models.StatusLineNode) error {
diff --git a/pkg/commands/models/status_line_node.go b/pkg/commands/models/status_line_node.go
index ce85cfec8..3962e0c42 100644
--- a/pkg/commands/models/status_line_node.go
+++ b/pkg/commands/models/status_line_node.go
@@ -13,24 +13,6 @@ type StatusLineNode struct {
 	Collapsed bool
 }
 
-func (s *StatusLineNode) GetShortStatus() string {
-	// need to see if any child has unstaged changes.
-	if s.IsLeaf() {
-		return s.File.ShortStatus
-	}
-
-	firstChar := " "
-	secondChar := " "
-	if s.GetHasStagedChanges() {
-		firstChar = "M"
-	}
-	if s.GetHasUnstagedChanges() {
-		secondChar = "M"
-	}
-
-	return firstChar + secondChar
-}
-
 func (s *StatusLineNode) GetHasUnstagedChanges() bool {
 	if s.IsLeaf() {
 		return s.File.HasUnstagedChanges
diff --git a/pkg/gui/presentation/files.go b/pkg/gui/presentation/files.go
index d64437b2c..900a5cabe 100644
--- a/pkg/gui/presentation/files.go
+++ b/pkg/gui/presentation/files.go
@@ -19,45 +19,53 @@ func GetFileListDisplayStrings(files []*models.File, diffName string, submoduleC
 
 // getFileDisplayStrings returns the display string of branch
 func getFileDisplayStrings(f *models.File, diffName string, submoduleConfigs []*models.SubmoduleConfig) []string {
-	output := GetStatusNodeLine(f.HasUnstagedChanges, f.ShortStatus, f.Name, diffName, submoduleConfigs, f)
+	output := GetStatusNodeLine(f.HasUnstagedChanges, f.HasStagedChanges, f.Name, diffName, submoduleConfigs, f)
 
 	return []string{output}
 }
 
-func GetStatusNodeLine(hasUnstagedChanges bool, shortStatus string, name string, diffName string, submoduleConfigs []*models.SubmoduleConfig, file *models.File) string {
+func GetStatusNodeLine(hasUnstagedChanges bool, hasStagedChanges bool, name string, diffName string, submoduleConfigs []*models.SubmoduleConfig, file *models.File) string {
 	// potentially inefficient to be instantiating these color
 	// objects with each render
 	red := color.New(color.FgRed)
 	green := color.New(color.FgGreen)
 	diffColor := color.New(theme.DiffTerminalColor)
+	partiallyModifiedColor := color.New(color.FgYellow)
 
 	var restColor *color.Color
 	if name == diffName {
 		restColor = diffColor
+	} else if file == nil && hasStagedChanges && hasUnstagedChanges {
+		restColor = partiallyModifiedColor
 	} else if hasUnstagedChanges {
 		restColor = red
 	} else {
 		restColor = green
 	}
 
-	// this is just making things look nice when the background attribute is 'reverse'
-	firstChar := shortStatus[0:1]
-	firstCharCl := green
-	if firstChar == "?" {
-		firstCharCl = red
-	} else if firstChar == " " {
-		firstCharCl = restColor
-	}
+	output := ""
+	if file != nil {
+		// this is just making things look nice when the background attribute is 'reverse'
+		firstChar := file.ShortStatus[0:1]
+		firstCharCl := green
+		if firstChar == "?" {
+			firstCharCl = red
+		} else if firstChar == " " {
+			firstCharCl = restColor
+		}
+
+		secondChar := file.ShortStatus[1:2]
+		secondCharCl := red
+		if secondChar == " " {
+			secondCharCl = restColor
+		}
 
-	secondChar := shortStatus[1:2]
-	secondCharCl := red
-	if secondChar == " " {
-		secondCharCl = restColor
+		output = firstCharCl.Sprint(firstChar)
+		output += secondCharCl.Sprint(secondChar)
+		output += " "
 	}
 
-	output := firstCharCl.Sprint(firstChar)
-	output += secondCharCl.Sprint(secondChar)
-	output += restColor.Sprintf(" %s", name)
+	output += restColor.Sprint(name)
 
 	if file != nil && file.IsSubmodule(submoduleConfigs) {
 		output += utils.ColoredString(" (submodule)", theme.DefaultTextColor)
diff --git a/pkg/gui/status_line_manager.go b/pkg/gui/status_line_manager.go
index 3da07ac5d..ba07c16b4 100644
--- a/pkg/gui/status_line_manager.go
+++ b/pkg/gui/status_line_manager.go
@@ -84,7 +84,7 @@ func (m *StatusLineManager) renderAux(s *models.StatusLineNode, prefix string, d
 	}
 
 	getLine := func() string {
-		return prefix + presentation.GetStatusNodeLine(s.GetHasUnstagedChanges(), s.GetShortStatus(), s.Name, diffName, submoduleConfigs, s.File)
+		return prefix + presentation.GetStatusNodeLine(s.GetHasUnstagedChanges(), s.GetHasStagedChanges(), s.Name, diffName, submoduleConfigs, s.File)
 	}
 
 	if s.IsLeaf() {