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

Code before:
		{
			Key:         opts.GetKey(opts.Config.Branches.SetUpstream),
			Handler:     self.checkSelected(self.setUpstream),
			Description: "(Un)set upstream",
			OpensMenu:   true,
		},
	}
}

// TODO: this is the same as SyncController.parseUpstream
func (self *BranchesController) parseUpstream(upstream string) (string, string, error) {
	var upstreamBranch, upstreamRemote string
	split := strings.Split(upstream, " ")
	if len(split) != 2 {
		return "", "", errors.New(self.c.Tr.InvalidUpstream)
	}

	upstreamRemote = split[0]
	upstreamBranch = split[1]

	return upstreamRemote, upstreamBranch, nil
}

func (self *BranchesController) promptForUpstream(currentBranch *models.Branch, onConfirm func(string) error) error {
	// TODO: how to get this from `gui`?
	// suggestedRemote := self.getSuggestedRemote()
	suggestedRemote := "origin"

	return self.c.Prompt(types.PromptOpts{
		Title:               self.c.Tr.EnterUpstream,
		InitialContent:      suggestedRemote + " " + currentBranch.Name,
		FindSuggestionsFunc: self.helpers.Suggestions.GetRemoteBranchesSuggestionsFunc(" "),
		HandleConfirm:       onConfirm,
	})
}

func (self *BranchesController) setUpstream(selectedBranch *models.Branch) error {
	return self.c.Menu(types.CreateMenuOptions{
		Title: "(Un)set upstream",
		Items: []*types.MenuItem{
			{
				DisplayStrings: []string{"unset upstream"},
				OnPress: func() error {
					return self.git.Branch.UnsetUpstream(selectedBranch.Name)
				},
				Key: 'u',
			},
			{
				DisplayStrings: []string{"set upstream"},
				OnPress: func() error {
					// TODO: this is the same as SyncController.promptForUpstream

Review comment (refactor):
Currently SyncController receives `getSuggestedRemote` from the gui package as an argument, but given that we need to use it here I'd move it into a controller helper. I'd add a helper called UpstreamHelper which contains GetSuggestedRemote, PromptForUpstream, and ParseUpstream, and then use those methods from here like `self.helpers.Upstream.PromptForUpstream`. I'm imagining something like this for the struct:
```go
type UpstreamHelper struct {
	c        *types.HelperCommon
	model    *types.Model
	suggestionsHelper *SuggestionsHelper
}
```
We'll need suggestionsHelper there so that we can call `self.suggestionsHelper.GetRemoteBranchesSuggestionsFunc`

What do you reckon?

Code after:
		{
			Key:         opts.GetKey(opts.Config.Branches.SetUpstream),
			Handler:     self.checkSelected(self.setUpstream),
			Description: self.c.Tr.LcSetUnsetUpstream,
			OpensMenu:   true,
		},
	}
}

func (self *BranchesController) setUpstream(selectedBranch *models.Branch) error {
	return self.c.Menu(types.CreateMenuOptions{
		Title: self.c.Tr.Actions.SetUnsetUpstream,
		Items: []*types.MenuItem{
			{
				DisplayStrings: []string{self.c.Tr.LcUnsetUpstream},
				OnPress: func() error {
					if err := self.git.Branch.UnsetUpstream(selectedBranch.Name); err != nil {
						return self.c.Error(err)
					}
					if err := self.c.Refresh(types.RefreshOptions{
						Mode: types.SYNC,
						Scope: []types.RefreshableView{
							types.BRANCHES,
							types.COMMITS,
						},
					}); err != nil {
						return self.c.Error(err)
					}
					return nil
				},
				Key: 'u',
			},
			{
				DisplayStrings: []string{self.c.Tr.LcSetUpstream},
				OnPress: func() error {
					return self.helpers.Upstream.PromptForUpstreamWithoutInitialContent(selectedBranch, func(upstream string) error {
						upstreamRemote, upstreamBranch, err := self.helpers.Upstream.ParseUpstream(upstream)
						if err != nil {
							return self.c.Error(err)
						}

						if err := self.git.Branch.SetUpstream(upstreamRemote, upstreamBranch, selectedBranch.Name); err != nil {
							return self.c.Error(err)
						}
						if err := self.c.Refresh(types.RefreshOptions{
							Mode: types.SYNC,
							Scope: []types.RefreshableView{
								types.BRANCHES,
								types.COMMITS,
							},
						}); err != nil {