Repository: jesseduffield/lazygit
File: pkg/gui/controllers/commits_files_controller.go (Go)

Code before:
	if err := self.c.Git().WorkingTree.CheckoutFile(self.context().GetRef().RefName(), node.GetPath()); err != nil {
		return self.c.Error(err)
	}

	return self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC})
}

func (self *CommitFilesController) discard(selectedNodes []*filetree.CommitFileNode) error {
	parentContext, ok := self.c.CurrentContext().GetParentContext()
	if !ok || parentContext.GetKey() != context.LOCAL_COMMITS_CONTEXT_KEY {
		return self.c.ErrorMsg(self.c.Tr.CanOnlyDiscardFromLocalCommits)
	}

	if ok, err := self.c.Helpers().PatchBuilding.ValidateNormalWorkingTreeState(); !ok {
		return err
	}

	removeFileRange := func() error {
		return self.c.WithWaitingStatus(self.c.Tr.RebasingStatus, func(gocui.Task) error {
			selectedNodes = normalisedSelectedCommitFileNodes(selectedNodes)

			return self.c.Confirm(types.ConfirmOpts{
				Title:  self.c.Tr.DiscardFileChangesTitle,
				Prompt: self.c.Tr.DiscardFileChangesPrompt,
				HandleConfirm: func() error {
					var filePaths []string

					// Reset the current patch if there is one.
					if self.c.Git().Patch.PatchBuilder.Active() {
						self.c.Git().Patch.PatchBuilder.Reset()
						if err := self.c.Refresh(types.RefreshOptions{Mode: types.BLOCK_UI}); err != nil {
							return err
						}
					}

					for _, node := range selectedNodes {
						err := node.ForEachFile(func(file *models.CommitFile) error {
							filePaths = append(filePaths, file.GetPath())
							return nil
						})
						if err != nil {
							return self.c.Error(err)
						}
					}

					err := self.c.Git().Rebase.DiscardOldFileChanges(self.c.Model().Commits, self.c.Contexts().LocalCommits.GetSelectedLineIdx(), filePaths)
					if err := self.c.Helpers().MergeAndRebase.CheckMergeOrRebase(err); err != nil {
						return err
					}

					return self.c.Refresh(types.RefreshOptions{Mode: types.BLOCK_UI})

Review comment (suggestion):
I need to build a list of all the individual files that are covered under `selectedNodes`. I'm feeling like there might be a better way to do this, but I didn't see a clean way to do it without just replicating the recursion that's already in `node.ForEachFile`.

Code after:
	if err := self.c.Git().WorkingTree.CheckoutFile(self.context().GetRef().RefName(), node.GetPath()); err != nil {
		return self.c.Error(err)
	}

	return self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC})
}

func (self *CommitFilesController) discard(selectedNodes []*filetree.CommitFileNode) error {
	parentContext, ok := self.c.CurrentContext().GetParentContext()
	if !ok || parentContext.GetKey() != context.LOCAL_COMMITS_CONTEXT_KEY {
		return self.c.ErrorMsg(self.c.Tr.CanOnlyDiscardFromLocalCommits)
	}

	if ok, err := self.c.Helpers().PatchBuilding.ValidateNormalWorkingTreeState(); !ok {
		return err
	}

	return self.c.Confirm(types.ConfirmOpts{
		Title:  self.c.Tr.DiscardFileChangesTitle,
		Prompt: self.c.Tr.DiscardFileChangesPrompt,
		HandleConfirm: func() error {
			return self.c.WithWaitingStatus(self.c.Tr.RebasingStatus, func(gocui.Task) error {
				var filePaths []string
				selectedNodes = normalisedSelectedCommitFileNodes(selectedNodes)

				// Reset the current patch if there is one.
				if self.c.Git().Patch.PatchBuilder.Active() {
					self.c.Git().Patch.PatchBuilder.Reset()
					if err := self.c.Refresh(types.RefreshOptions{Mode: types.BLOCK_UI}); err != nil {
						return err
					}
				}

				for _, node := range selectedNodes {
					err := node.ForEachFile(func(file *models.CommitFile) error {
						filePaths = append(filePaths, file.GetPath())
						return nil
					})
					if err != nil {
						return self.c.Error(err)
					}
				}

				err := self.c.Git().Rebase.DiscardOldFileChanges(self.c.Model().Commits, self.c.Contexts().LocalCommits.GetSelectedLineIdx(), filePaths)
				if err := self.c.Helpers().MergeAndRebase.CheckMergeOrRebase(err); err != nil {
					return err
				}

				if self.context().RangeSelectEnabled() {
					self.context().GetList().CancelRangeSelect()
				}