Implement state filtering for commit files

3 files changed, 81 insertions, 2 deletions

diff --git a/pkg/gui/files_panel.go b/pkg/gui/files_panel.go
index f652c3f05..ac9b117eb 100644
--- a/pkg/gui/files_panel.go
+++ b/pkg/gui/files_panel.go
@@ -473,6 +473,39 @@ func (gui *Gui) handleCommitEditorPress() error {
 	)
 }
 
+func (gui *Gui) handleStatusFilterPressed() error {
+	menuItems := []*menuItem{}
+
+	menuItems = append(menuItems, &menuItem{
+		displayString: "Show only staged files",
+		onPress: func() error {
+			return gui.setStatusFiltering(filetree.DisplayStaged)
+		},
+	})
+
+	menuItems = append(menuItems, &menuItem{
+		displayString: "Show only unstaged files",
+		onPress: func() error {
+			return gui.setStatusFiltering(filetree.DisplayUnstaged)
+		},
+	})
+
+	menuItems = append(menuItems, &menuItem{
+		displayString: "Reset Filter",
+		onPress: func() error {
+			return gui.setStatusFiltering(filetree.DisplayAll)
+		},
+	})
+
+	return gui.createMenu(gui.Tr.FilteringMenuTitle, menuItems, createMenuOptions{showCancel: false})
+}
+
+func (gui *Gui) setStatusFiltering(filter filetree.FileManagerDisplayFilter) error {
+	state := gui.State
+	state.FileManager.SetDisplayFilter(filter)
+	return gui.handleRefreshFiles()
+}
+
 func (gui *Gui) editFile(filename string) error {
 	return gui.editFileAtLine(filename, 1)
 }
diff --git a/pkg/gui/filetree/file_manager.go b/pkg/gui/filetree/file_manager.go
index 825ee0b01..699e20332 100644
--- a/pkg/gui/filetree/file_manager.go
+++ b/pkg/gui/filetree/file_manager.go
@@ -8,11 +8,20 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+type FileManagerDisplayFilter int
+
+const (
+	DisplayAll FileManagerDisplayFilter = iota
+	DisplayStaged
+	DisplayUnstaged
+)
+
 type FileManager struct {
 	files          []*models.File
 	tree           *FileNode
 	showTree       bool
 	log            *logrus.Entry
+	filter         FileManagerDisplayFilter
 	collapsedPaths CollapsedPaths
 	sync.RWMutex
 }
@@ -22,6 +31,7 @@ func NewFileManager(files []*models.File, log *logrus.Entry, showTree bool) *Fil
 		files:          files,
 		log:            log,
 		showTree:       showTree,
+		filter:         DisplayAll,
 		collapsedPaths: CollapsedPaths{},
 		RWMutex:        sync.RWMutex{},
 	}
@@ -35,6 +45,35 @@ func (m *FileManager) ExpandToPath(path string) {
 	m.collapsedPaths.ExpandToPath(path)
 }
 
+func (m *FileManager) GetFilesForDisplay() []*models.File {
+	files := m.files
+	if m.filter == DisplayAll {
+		return files
+	}
+
+	result := make([]*models.File, 0)
+	if m.filter == DisplayStaged {
+		for _, file := range files {
+			if file.HasStagedChanges {
+				result = append(result, file)
+			}
+		}
+	} else {
+		for _, file := range files {
+			if !file.HasStagedChanges {
+				result = append(result, file)
+			}
+		}
+	}
+
+	return result
+}
+
+func (m *FileManager) SetDisplayFilter(filter FileManagerDisplayFilter) {
+	m.filter = filter
+	m.SetTree()
+}
+
 func (m *FileManager) ToggleShowTree() {
 	m.showTree = !m.showTree
 	m.SetTree()
@@ -73,10 +112,11 @@ func (m *FileManager) SetFiles(files []*models.File) {
 }
 
 func (m *FileManager) SetTree() {
+	filesForDisplay := m.GetFilesForDisplay()
 	if m.showTree {
-		m.tree = BuildTreeFromFiles(m.files)
+		m.tree = BuildTreeFromFiles(filesForDisplay)
 	} else {
-		m.tree = BuildFlatTreeFromFiles(m.files)
+		m.tree = BuildFlatTreeFromFiles(filesForDisplay)
 	}
 }
 
diff --git a/pkg/gui/keybindings.go b/pkg/gui/keybindings.go
index 4b20bab95..4d593dbd0 100644
--- a/pkg/gui/keybindings.go
+++ b/pkg/gui/keybindings.go
@@ -383,6 +383,12 @@ func (gui *Gui) GetInitialKeybindings() []*Binding {
 		},
 		{
 			ViewName:    "files",
+			Key:         gui.getKey("<c-b>"),
+			Handler:     gui.handleStatusFilterPressed,
+			Description: gui.Tr.LcCommitFileFilter,
+		},
+		{
+			ViewName:    "files",
 			Contexts:    []string{string(FILES_CONTEXT_KEY)},
 			Key:         gui.getKey(config.Files.CommitChanges),
 			Handler:     gui.handleCommitPress,