diff --git a/pkg/gui/controllers/files_controller.go b/pkg/gui/controllers/files_controller.go index 5f4a3bae0..e37127181 100644 --- a/pkg/gui/controllers/files_controller.go +++ b/pkg/gui/controllers/files_controller.go @@ -6,6 +6,7 @@ import ( "path/filepath" "strings" + "github.com/jesseduffield/generics/set" "github.com/jesseduffield/gocui" "github.com/jesseduffield/lazygit/pkg/commands/git_commands" "github.com/jesseduffield/lazygit/pkg/commands/models" @@ -1194,13 +1195,36 @@ func filterNodesHaveUnstagedChanges(nodes []*filetree.FileNode) []*filetree.File }) } +func findSubmoduleNode(nodes []*filetree.FileNode, submodules []*models.SubmoduleConfig) *models.File { + for _, node := range nodes { + submoduleNode := node.FindFirstFileBy(func(f *models.File) bool { + return f.IsSubmodule(submodules) + }) + if submoduleNode != nil { + return submoduleNode + } + } + return nil +} + func (self *FilesController) canRemove(selectedNodes []*filetree.FileNode) *types.DisabledReason { + // Return disabled if the selection contains multiple changed items and includes a submodule change. submodules := self.c.Model().Submodules - submoduleCount := lo.CountBy(selectedNodes, func(node *filetree.FileNode) bool { - return node.File != nil && node.File.IsSubmodule(submodules) - }) - if submoduleCount > 0 && len(selectedNodes) > 1 { - return &types.DisabledReason{Text: self.c.Tr.RangeSelectNotSupportedForSubmodules} + hasFiles := false + uniqueSelectedSubmodules := set.New[*models.SubmoduleConfig]() + + for _, node := range selectedNodes { + _ = node.ForEachFile(func(f *models.File) error { + if submodule := f.SubmoduleConfig(submodules); submodule != nil { + uniqueSelectedSubmodules.Add(submodule) + } else { + hasFiles = true + } + return nil + }) + if uniqueSelectedSubmodules.Len() > 0 && (hasFiles || uniqueSelectedSubmodules.Len() > 1) { + return &types.DisabledReason{Text: self.c.Tr.RangeSelectNotSupportedForSubmodules} + } } return nil @@ -1209,11 +1233,13 @@ func (self *FilesController) canRemove(selectedNodes []*filetree.FileNode) *type func (self *FilesController) remove(selectedNodes []*filetree.FileNode) error { submodules := self.c.Model().Submodules + selectedNodes = normalisedSelectedNodes(selectedNodes) + // If we have one submodule then we must only have one submodule or `canRemove` would have // returned an error - firstNode := selectedNodes[0] - if firstNode.File != nil && firstNode.File.IsSubmodule(submodules) { - submodule := firstNode.File.SubmoduleConfig(submodules) + submoduleNode := findSubmoduleNode(selectedNodes, submodules) + if submoduleNode != nil { + submodule := submoduleNode.SubmoduleConfig(submodules) menuItems := []*types.MenuItem{ { @@ -1224,11 +1250,9 @@ func (self *FilesController) remove(selectedNodes []*filetree.FileNode) error { }, } - return self.c.Menu(types.CreateMenuOptions{Title: firstNode.GetPath(), Items: menuItems}) + return self.c.Menu(types.CreateMenuOptions{Title: submoduleNode.GetPath(), Items: menuItems}) } - selectedNodes = normalisedSelectedNodes(selectedNodes) - discardAllChangesItem := types.MenuItem{ Label: self.c.Tr.DiscardAllChanges, OnPress: func() error { diff --git a/pkg/gui/filetree/node.go b/pkg/gui/filetree/node.go index ce3bac566..97d5232b5 100644 --- a/pkg/gui/filetree/node.go +++ b/pkg/gui/filetree/node.go @@ -169,6 +169,22 @@ func (self *Node[T]) EveryFile(predicate func(*T) bool) bool { return true } +func (self *Node[T]) FindFirstFileBy(predicate func(*T) bool) *T { + if self.IsFile() { + if predicate(self.File) { + return self.File + } + } else { + for _, child := range self.Children { + if file := child.FindFirstFileBy(predicate); file != nil { + return file + } + } + } + + return nil +} + func (self *Node[T]) Flatten(collapsedPaths *CollapsedPaths) []*Node[T] { result := []*Node[T]{self}