[git] IJPL-72576 Add option to set upstream in push dialog

GitOrigin-RevId: 7c50f0d3fe29bec2113cd5073df00224cdc38d9d

platform/dvcs-api/api-dump.txt

@@ -88,12 +88,14 @@ a:com.intellij.dvcs.push.PushTargetPanel
 - javax.swing.JPanel
 - <init>():V
 - a:addTargetEditorListener(com.intellij.dvcs.push.ui.PushTargetEditorListener):V
+- editingStarted():V
 - a:fireOnCancel():V
 - a:fireOnChange():V
 - forceUpdateEditableUiModel(java.lang.String):V
 - a:getValue():com.intellij.dvcs.push.PushTarget
 - a:render(com.intellij.ui.ColoredTreeCellRenderer,Z,Z,java.lang.String):V
 - a:setFireOnChangeAction(java.lang.Runnable):V
+- showSourceWhenEditing():Z
 - a:verify():com.intellij.openapi.ui.ValidationInfo
 a:com.intellij.dvcs.push.Pusher
 - <init>():V

platform/dvcs-api/src/com/intellij/dvcs/push/PushTargetPanel.java

@@ -37,6 +37,8 @@ public abstract class PushTargetPanel<T extends PushTarget> extends JPanel {
   @Nullable
   abstract public T getValue();
 
+  public void editingStarted() { }
+
   public abstract void fireOnCancel();
 
   public abstract void fireOnChange();
@@ -53,4 +55,8 @@ public abstract class PushTargetPanel<T extends PushTarget> extends JPanel {
 
   public void forceUpdateEditableUiModel(@NotNull String forcedText) {
   }
+
+  public boolean showSourceWhenEditing() {
+    return true;
+  }
 }

platform/dvcs-impl/src/com/intellij/dvcs/push/ui/PushLog.java

@@ -83,7 +83,6 @@ public final class PushLog extends JPanel implements Disposable, UiDataProvider
     treeModel.nodeStructureChanged(root);
     myTreeCellRenderer = new MyTreeCellRenderer();
     myTree = new CheckboxTree(myTreeCellRenderer, root) {
-
       @Override
       protected boolean shouldShowBusyIconIfNeeded() {
         return true;
@@ -745,6 +744,16 @@ public final class PushLog extends JPanel implements Disposable, UiDataProvider
         }
       };
     }
+
+    @Override
+    protected boolean startEditing(TreePath path, MouseEvent event) {
+      boolean editingStarted = super.startEditing(path, event);
+      if (editingStarted && myTree.getCellEditor() instanceof MyTreeCellEditor editor) {
+        editor.myValue.getTargetPanel().editingStarted();
+      }
+
+      return editingStarted;
+    }
   }
 
   private static class MyTreeViewPort extends JBViewport {

platform/dvcs-impl/src/com/intellij/dvcs/push/ui/RepositoryWithBranchPanel.java

@@ -116,7 +116,10 @@ public class RepositoryWithBranchPanel<T extends PushTarget> extends NonOpaquePa
     Rectangle bounds = tree.getPathBounds(tree.getPathForRow(row));
     invalidate();
     myTextRenderer.getTreeCellRendererComponent(tree, value, selected, expanded, leaf, row, hasFocus);
-    if (!(value instanceof SingleRepositoryNode)) {
+    if (value instanceof SingleRepositoryNode) {
+      myTextRenderer.setIpad(JBUI.insetsLeft(10));
+      myRepositoryCheckbox.setVisible(false);
+    } else {
       RepositoryNode node = (RepositoryNode)value;
       myRepositoryCheckbox.setSelected(node.isChecked());
       myRepositoryCheckbox.setVisible(true);
@@ -124,13 +127,14 @@ public class RepositoryWithBranchPanel<T extends PushTarget> extends NonOpaquePa
       myTextRenderer.append(getRepositoryName(), SimpleTextAttributes.GRAY_ATTRIBUTES);
       myTextRenderer.appendTextPadding(120);
     }
-    else {
+
-      myTextRenderer.setIpad(JBUI.insetsLeft(10));
+    if (myDestPushTargetPanelComponent.showSourceWhenEditing()) {
-      myRepositoryCheckbox.setVisible(false);
+      if (value instanceof SingleRepositoryNode) {
-      myTextRenderer.append(" ");
+        myTextRenderer.append(" ");
+      }
+      myTextRenderer.append(getSourceName(), SimpleTextAttributes.REGULAR_ATTRIBUTES);
+      myTextRenderer.append(getArrow(), SimpleTextAttributes.REGULAR_ATTRIBUTES);
     }
-    myTextRenderer.append(getSourceName(), SimpleTextAttributes.REGULAR_ATTRIBUTES);
-    myTextRenderer.append(getArrow(), SimpleTextAttributes.REGULAR_ATTRIBUTES);
     if (bounds != null) {
       setPreferredSize(new Dimension(tree.getVisibleRect().width - bounds.x, bounds.height));
     }

plugins/git4idea/resources/messages/GitBundle.properties

@@ -776,6 +776,9 @@ push.dialog.target.panel.adding.remote=Adding Remote\u2026
 push.dialog.target.panel.can.t.push=Cannot push
 push.dialog.target.panel.empty.repository=Empty repository
 push.dialog.target.panel.new=New
+push.dialog.target.panel.upstream.checkbox=Set upstream
+push.dialog.target.panel.upstream.label=Upstream
+push.dialog.target.panel.new.and.upstream=New \\& Upstream
 push.dialog.preview.commits.before.push=For Commit and Push to non-protected branches, preview commits before push
 push.dialog.prohibited.branch.configurable.path={0} | Version Control | Git
 push.local.history.system.label.after=After push

plugins/git4idea/src/git4idea/GitOperationsCollector.kt

@@ -8,12 +8,13 @@ import com.intellij.internal.statistic.service.fus.collectors.CounterUsagesColle
 import com.intellij.openapi.project.Project
 import git4idea.commands.GitCommandResult
 import git4idea.push.GitPushRepoResult
+import git4idea.push.GitPushTarget
 import git4idea.push.GitPushTargetType
 
 object GitOperationsCollector : CounterUsagesCollector() {
   override fun getGroup(): EventLogGroup = GROUP
 
-  private val GROUP: EventLogGroup = EventLogGroup("git.operations", 4)
+  private val GROUP: EventLogGroup = EventLogGroup("git.operations", 5)
 
   internal val UPDATE_FORCE_PUSHED_BRANCH_ACTIVITY = GROUP.registerIdeActivity("update.force.pushed")
 
@@ -22,11 +23,15 @@ object GitOperationsCollector : CounterUsagesCollector() {
   private val PUSHED_COMMITS_COUNT = EventFields.RoundedInt("pushed_commits_count")
   private val PUSH_RESULT = EventFields.Enum<GitPushRepoResult.Type>("push_result")
   private val TARGET_TYPE  = EventFields.Enum<GitPushTargetType>("push_target_type")
+  private val SET_UPSTREAM  = EventFields.Boolean("push_set_upsteram")
+  private val PUSH_TO_NEW_BRANCH = EventFields.Boolean("push_new_branch")
   private val PUSH_ACTIVITY = GROUP.registerIdeActivity("push",
                                                         finishEventAdditionalFields = arrayOf(PUSHED_COMMITS_COUNT,
                                                                                               PUSH_RESULT,
                                                                                               IS_AUTHENTICATION_FAILED,
-                                                                                              TARGET_TYPE))
+                                                                                              TARGET_TYPE,
+                                                                                              SET_UPSTREAM,
+                                                                                              PUSH_TO_NEW_BRANCH))
 
   private val EXPECTED_COMMITS_NUMBER = EventFields.Int("expected_commits_number")
   private val ACTUAL_COMMITS_NUMBER = EventFields.Int("actual_commits_number")
@@ -42,12 +47,22 @@ object GitOperationsCollector : CounterUsagesCollector() {
   }
 
   @JvmStatic
-  fun endLogPush(activity: StructuredIdeActivity, commandResult: GitCommandResult?, pushRepoResult: GitPushRepoResult?, targetType: GitPushTargetType?) {
+  fun endLogPush(
+    activity: StructuredIdeActivity,
+    commandResult: GitCommandResult?,
+    pushRepoResult: GitPushRepoResult?,
+    targetType: GitPushTargetType?,
+    newBranchCreated: Boolean,
+    setUpstream: Boolean,
+  ) {
     activity.finished {
-      listOfNotNull(pushRepoResult?.let { PUSHED_COMMITS_COUNT with it.numberOfPushedCommits },
+      listOfNotNull(
-                    pushRepoResult?.let { PUSH_RESULT with it.type },
+        pushRepoResult?.let { PUSHED_COMMITS_COUNT with it.numberOfPushedCommits },
-                    commandResult?.let { IS_AUTHENTICATION_FAILED with it.isAuthenticationFailed },
+        pushRepoResult?.let { PUSH_RESULT with it.type },
-                    targetType?.let { TARGET_TYPE with it }
+        commandResult?.let { IS_AUTHENTICATION_FAILED with it.isAuthenticationFailed },
+        targetType?.let { TARGET_TYPE with it },
+        PUSH_TO_NEW_BRANCH with newBranchCreated,
+        SET_UPSTREAM with setUpstream,
       )
     }
   }

plugins/git4idea/src/git4idea/push/GitPushOperation.java

@@ -277,9 +277,9 @@ public class GitPushOperation {
       GitPushRepoResult repoResult = null;
 
       StructuredIdeActivity pushActivity = GitOperationsCollector.startLogPush(repository.getProject());
-      GitPushTargetType targetType = getPushTargetType(repository, spec);
+      boolean setUpstream = spec.getTarget().shouldSetUpstream(spec.getSource(), repository);
       try {
-        resultWithOutput = doPush(repository, spec);
+        resultWithOutput = doPush(repository, spec, setUpstream);
         LOG.debug("Pushed to " + DvcsUtil.getShortRepositoryName(repository) + ": " + resultWithOutput);
 
         GitPushSource pushSource = spec.getSource();
@@ -308,7 +308,14 @@ public class GitPushOperation {
         }
       }
       finally {
-        GitOperationsCollector.endLogPush(pushActivity, resultWithOutput != null ? resultWithOutput.resultOutput : null, repoResult, targetType);
+        GitOperationsCollector.endLogPush(
+          pushActivity,
+          resultWithOutput != null ? resultWithOutput.resultOutput : null,
+          repoResult,
+          getPushTargetType(repository, spec),
+          spec.getTarget().isNewBranchCreated(),
+          setUpstream
+        );
       }
 
       LOG.debug("Converted result: " + repoResult);
@@ -392,20 +399,14 @@ public class GitPushOperation {
     }
   }
 
-  private @NotNull ResultWithOutput doPush(@NotNull GitRepository repository, @NotNull PushSpec<GitPushSource, GitPushTarget> pushSpec) {
+  private @NotNull ResultWithOutput doPush(@NotNull GitRepository repository, @NotNull PushSpec<GitPushSource, GitPushTarget> pushSpec, boolean setUpstream) {
-    GitPushSource pushSource = pushSpec.getSource();
     GitPushTarget pushTarget = pushSpec.getTarget();
-    GitLocalBranch sourceBranch = pushSource.getBranch();
     GitRemoteBranch targetBranch = pushTarget.getBranch();
 
     GitLineHandlerListener progressListener = GitStandardProgressAnalyzer.createListener(myProgressIndicator);
-    boolean setUpstream = sourceBranch != null &&
-                          pushTarget.isNewBranchCreated() &&
-                          pushSource.isBranchRef() &&
-                          !branchTrackingInfoIsSet(repository, sourceBranch);
     String tagMode = myTagMode == null ? null : myTagMode.getArgument();
 
-    String spec = createPushSpec(pushSource, pushTarget, setUpstream);
+    String spec = createPushSpec(pushSpec.getSource(), pushTarget, setUpstream);
     GitRemote remote = targetBranch.getRemote();
 
     List<GitPushParams.ForceWithLease> forceWithLease = emptyList();

plugins/git4idea/src/git4idea/push/GitPushTarget.java

@@ -34,6 +34,8 @@ public class GitPushTarget implements PushTarget {
   private final boolean myPushingToSpecialRef;
   private final @Nullable GitPushTargetType myTargetType;
 
+  private boolean shouldSetNewUpstream;
+
   public GitPushTarget(@NotNull GitRemoteBranch remoteBranch, boolean isNewBranchCreated) {
     this(remoteBranch, isNewBranchCreated, false, null);
   }
@@ -76,6 +78,18 @@ public class GitPushTarget implements PushTarget {
     return myPushingToSpecialRef;
   }
 
+  public void shouldSetNewUpstream(boolean value) {
+    shouldSetNewUpstream = value;
+  }
+
+  public boolean shouldSetUpstream(@NotNull GitPushSource pushSource, @NotNull GitRepository repository) {
+    GitLocalBranch sourceBranch = pushSource.getBranch();
+
+    return sourceBranch != null &&
+           pushSource.isBranchRef() &&
+           ((isNewBranchCreated() && !branchTrackingInfoIsSet(repository, sourceBranch)) || shouldSetNewUpstream);
+  }
+
   public @Nullable GitPushTargetType getTargetType() {
     return myTargetType;
   }
@@ -146,6 +160,10 @@ public class GitPushTarget implements PushTarget {
     return getDefaultOrFirstRemote(repository.getRemotes());
   }
 
+  private static boolean branchTrackingInfoIsSet(@NotNull GitRepository repository, final @NotNull GitLocalBranch source) {
+    return ContainerUtil.exists(repository.getBranchTrackInfos(), info -> info.getLocalBranch().equals(source));
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;
@@ -153,6 +171,7 @@ public class GitPushTarget implements PushTarget {
 
     if (myIsNewBranchCreated != target.myIsNewBranchCreated) return false;
     if (!myRemoteBranch.equals(target.myRemoteBranch)) return false;
+    if (shouldSetNewUpstream != target.shouldSetNewUpstream) return false;
 
     return true;
   }
@@ -161,6 +180,7 @@ public class GitPushTarget implements PushTarget {
   public int hashCode() {
     int result = myRemoteBranch.hashCode();
     result = 31 * result + (myIsNewBranchCreated ? 1 : 0);
+    result = 31 * result + (shouldSetNewUpstream ? 1 : 0);
     return result;
   }
 

plugins/git4idea/src/git4idea/push/GitPushTargetPanel.java

@@ -21,7 +21,7 @@ import com.intellij.openapi.util.NlsSafe;
 import com.intellij.openapi.wm.IdeFocusManager;
 import com.intellij.ui.*;
 import com.intellij.ui.awt.RelativePoint;
-import com.intellij.ui.components.JBLabel;
+import com.intellij.ui.components.JBCheckBox;
 import com.intellij.ui.popup.list.ListPopupImpl;
 import com.intellij.ui.scale.JBUIScale;
 import com.intellij.util.containers.ContainerUtil;
@@ -37,11 +37,11 @@ import git4idea.i18n.GitBundle;
 import git4idea.remote.GitDefineRemoteDialog;
 import git4idea.repo.GitRemote;
 import git4idea.repo.GitRepository;
+import git4idea.validators.GitRefNameValidator;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import javax.swing.*;
 import javax.swing.tree.DefaultMutableTreeNode;
 import java.awt.*;
 import java.awt.event.*;
@@ -58,14 +58,6 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
 
   private static final Comparator<GitRemoteBranch> REMOTE_BRANCH_COMPARATOR = new MyRemoteBranchComparator();
   private static final String SEPARATOR = " : ";
-  private static final Color NEW_BRANCH_LABEL_FG = new JBColor(0x00b53d, 0x6ba65d);
-  private static final Color NEW_BRANCH_LABEL_SELECTION_FG = UIUtil.getTreeSelectionForeground();
-  private static final Color NEW_BRANCH_LABEL_BG = new JBColor(0xebfcf1, 0x313b32);
-  private static final Color NEW_BRANCH_LABEL_SELECTION_BG =
-    new JBColor(ColorUtil.toAlpha(NEW_BRANCH_LABEL_SELECTION_FG, 20), ColorUtil.toAlpha(NEW_BRANCH_LABEL_SELECTION_FG, 30));
-  private static final RelativeFont NEW_BRANCH_LABEL_FONT = RelativeFont.TINY.small();
-  private static final TextIcon NEW_BRANCH_LABEL =
-    new TextIcon(GitBundle.message("push.dialog.target.panel.new"), NEW_BRANCH_LABEL_FG, NEW_BRANCH_LABEL_BG, 0);
 
   private final @NotNull GitPushSupport myPushSupport;
   private final @NotNull GitRepository myRepository;
@@ -76,6 +68,7 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
   private final @NotNull PushTargetTextField myTargetEditor;
   private final @NotNull VcsLinkedTextComponent myRemoteRenderer;
   private final @NotNull Project myProject;
+  private final @Nullable SetUpstreamCheckbox myUpstreamCheckbox;
 
   private @Nullable GitPushTarget myCurrentTarget;
   private @Nullable @Nls String myError;
@@ -99,6 +92,7 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
 
     myTargetRenderer = new VcsEditableTextComponent("", null);
     myTargetEditor = new PushTargetTextField(repository.getProject(), getTargetNames(myRepository), "");
+
     myRemoteRenderer = new VcsLinkedTextComponent("", new VcsLinkListener() {
       @Override
       public void hyperlinkActivated(@NotNull DefaultMutableTreeNode sourceNode, @NotNull MouseEvent event) {
@@ -114,16 +108,28 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
       }
     });
 
-    setLayout(new BorderLayout());
     setOpaque(false);
-    JPanel remoteAndSeparator = new JPanel(new BorderLayout());
-    remoteAndSeparator.setOpaque(false);
-    remoteAndSeparator.add(myRemoteRenderer, BorderLayout.CENTER);
-    remoteAndSeparator.add(new JBLabel(SEPARATOR), BorderLayout.EAST);
 
-    add(remoteAndSeparator, BorderLayout.WEST);
+    setLayout(new BorderLayout());
     add(myTargetEditor, BorderLayout.CENTER);
 
+    if (source instanceof GitPushSource.OnBranch && defaultTarget != null &&
+        // "Set upstream" checkbox isn't shown if there is no existing tracking branch
+        defaultTarget.getTargetType() == GitPushTargetType.TRACKING_BRANCH && !defaultTarget.isNewBranchCreated()
+    ) {
+      myUpstreamCheckbox = new SetUpstreamCheckbox(defaultTarget.getBranch().getNameForRemoteOperations());
+      myTargetEditor.addDocumentListener(new DocumentListener() {
+        @Override
+        public void documentChanged(@NotNull DocumentEvent event) {
+          myUpstreamCheckbox.setVisible(myTargetEditor.getText());
+        }
+      });
+      add(myUpstreamCheckbox, BorderLayout.EAST);
+    }
+    else {
+      myUpstreamCheckbox = null;
+    }
+
     updateComponents(defaultTarget);
 
     setFocusCycleRoot(true);
@@ -165,6 +171,10 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
     myTargetEditor.setText(initialBranch);
     myRemoteRenderer.updateLinkText(noRemotes ? GitBundle.message("push.dialog.target.panel.define.remote") : initialRemote);
 
+    if (myUpstreamCheckbox != null) {
+      myUpstreamCheckbox.setVisible(initialBranch);
+    }
+
     myTargetEditor.setVisible(!noRemotes);
   }
 
@@ -299,14 +309,19 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
         myTargetRenderer.setSelected(isSelected);
         myTargetRenderer.setTransparent(!isActive);
         myTargetRenderer.render(renderer);
-        if (newRemoteBranch) {
+        boolean newUpstream = myUpstreamCheckbox != null &&
+                              target != null &&
+                              myUpstreamCheckbox.isSelected() &&
+                              !myUpstreamCheckbox.isDefaultUpstream(target.getBranch().getNameForRemoteOperations());
+        if (newRemoteBranch || newUpstream) {
           renderer.setIconOnTheRight(true);
-          NEW_BRANCH_LABEL.setInsets(JBUI.insets(2));
+        }
-          NEW_BRANCH_LABEL.setRound(JBUIScale.scale(4));
+        if (newRemoteBranch && newUpstream) {
-          NEW_BRANCH_LABEL.setFont(NEW_BRANCH_LABEL_FONT.derive(renderer.getFont()));
+          renderer.setIcon(BranchLabels.getNewAndUpstreamBranchLabel(renderer.getFont(), isSelected));
-          NEW_BRANCH_LABEL.setForeground(isSelected ? NEW_BRANCH_LABEL_SELECTION_FG : NEW_BRANCH_LABEL_FG);
+        } else if (newRemoteBranch) {
-          NEW_BRANCH_LABEL.setBackground(isSelected ? NEW_BRANCH_LABEL_SELECTION_BG : NEW_BRANCH_LABEL_BG);
+          renderer.setIcon(BranchLabels.getNewBranchLabel(renderer.getFont(), isSelected));
-          renderer.setIcon(NEW_BRANCH_LABEL);
+        } else if (newUpstream) {
+          renderer.setIcon(BranchLabels.getUpstreamBranchLabel(renderer.getFont(), isSelected));
         }
       }
     }
@@ -321,19 +336,41 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
     return (target != null ? target.getBranch().getNameForRemoteOperations() : "");
   }
 
+  @Override
+  public void editingStarted() {
+    if (myUpstreamCheckbox != null) {
+      // Checkbox should be explicitly enabled and disabled when toggling editing, as click
+      // to start editing can change checkbox state
+      // See BasicTreeUi#startEditing for details
+      myUpstreamCheckbox.setVisible(myTargetEditor.getText());
+      myUpstreamCheckbox.setEnabled(true);
+    }
+  }
+
   @Override
   public void fireOnCancel() {
+    if (myUpstreamCheckbox != null) {
+      myUpstreamCheckbox.setEnabled(false);
+    }
+
     myTargetEditor.setText(getTextFieldText(myCurrentTarget));
   }
 
   @Override
   public void fireOnChange() {
+    if (myUpstreamCheckbox != null) {
+      myUpstreamCheckbox.setEnabled(false);
+    }
+
     //any changes are senselessly if no remotes
     if (myError != null || myRepository.getRemotes().isEmpty()) return;
     String remoteName = myRemoteRenderer.getText();
     String branchName = myTargetEditor.getText();
     try {
       GitPushTarget target = GitPushTarget.parse(myRepository, remoteName, branchName);
+      if (myUpstreamCheckbox != null) {
+        target.shouldSetNewUpstream(myUpstreamCheckbox.isVisible() && myUpstreamCheckbox.isSelected());
+      }
       if (!target.equals(myCurrentTarget)) {
         myCurrentTarget = target;
         myTargetRenderer.updateLinkText(branchName);
@@ -435,6 +472,11 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
     }
   }
 
+  @Override
+  public boolean showSourceWhenEditing() {
+    return false;
+  }
+
   private static final class PopupItem {
     static final PopupItem DEFINE_REMOTE = new PopupItem(null);
 
@@ -486,4 +528,57 @@ public class GitPushTargetPanel extends PushTargetPanel<GitPushTarget> {
       return aComponent;
     }
   }
+
+  private static class SetUpstreamCheckbox extends JBCheckBox {
+    private final String upstreamBranchName;
+
+    SetUpstreamCheckbox(String upstreamBranchName) {
+      super(GitBundle.message("push.dialog.target.panel.upstream.checkbox"), false);
+
+      this.upstreamBranchName = upstreamBranchName;
+      setBorder(JBUI.Borders.empty(0, 5, 0, 10));
+      setOpaque(false);
+      setFocusable(false);
+    }
+
+    public void setVisible(String targetName) {
+      boolean valid = GitRefNameValidator.getInstance().checkInput(targetName);
+      setVisible(valid && !isDefaultUpstream(targetName));
+    }
+
+    public boolean isDefaultUpstream(String targetName) {
+      return upstreamBranchName.equals(targetName);
+    }
+  }
+
+  private static final class BranchLabels {
+    private static final Color LABEL_FG = new JBColor(0x00b53d, 0x6ba65d);
+    private static final Color LABEL_SELECTION_FG = UIUtil.getTreeSelectionForeground();
+    private static final Color LABEL_BG = new JBColor(0xebfcf1, 0x313b32);
+    private static final Color LABEL_SELECTION_BG =
+      new JBColor(ColorUtil.toAlpha(LABEL_SELECTION_FG, 20), ColorUtil.toAlpha(LABEL_SELECTION_FG, 30));
+    private static final RelativeFont LABEL_FONT = RelativeFont.TINY.small();
+
+    public static TextIcon getNewBranchLabel(Font font, boolean selected) {
+      return getLabel(GitBundle.message("push.dialog.target.panel.new"), font, selected);
+    }
+
+    public static TextIcon getUpstreamBranchLabel(Font font, boolean selected) {
+      return getLabel(GitBundle.message("push.dialog.target.panel.upstream.label"), font, selected);
+    }
+
+    public static TextIcon getNewAndUpstreamBranchLabel(Font font, boolean selected) {
+      return getLabel(GitBundle.message("push.dialog.target.panel.new.and.upstream"), font, selected);
+    }
+
+    private static TextIcon getLabel(String text, Font font, boolean selected) {
+      TextIcon label = new TextIcon(text, LABEL_FG, LABEL_BG, 0);
+      label.setInsets(JBUI.insets(2));
+      label.setRound(JBUIScale.scale(4));
+      label.setFont(LABEL_FONT.derive(font));
+      label.setForeground(selected ? LABEL_SELECTION_FG : LABEL_FG);
+      label.setBackground(selected ? LABEL_SELECTION_BG : LABEL_BG);
+      return label;
+    }
+  }
 }

plugins/git4idea/tests/git4idea/push/GitPushOperationSingleRepoTest.kt

@@ -29,7 +29,6 @@ import org.junit.Before
 import java.io.File
 import java.nio.file.Path
 import java.util.Collections.singletonMap
-import kotlin.Throws
 
 class GitPushOperationSingleRepoTest : GitPushOperationBaseTest() {
   private lateinit var repository: GitRepository
@@ -448,6 +447,13 @@ class GitPushOperationSingleRepoTest : GitPushOperationBaseTest() {
     assertEmpty(pushedTags)
   }
 
+  fun `test push with setting upstream`() {
+    push("master", "origin/feature", canChangeUpstream = true)
+    assertUpstream("master", "origin", "feature")
+    push("master", "origin/feature-1", canChangeUpstream = true)
+    assertUpstream("master", "origin", "feature-1")
+  }
+
   fun `test skip pre push hook`() {
     assumeTrue("Not testing: pre-push hooks are not supported in ${vcs.version}", GitVersionSpecialty.PRE_PUSH_HOOK.existsIn(vcs.version))
 
@@ -494,12 +500,12 @@ class GitPushOperationSingleRepoTest : GitPushOperationBaseTest() {
     makeCommit("file.txt")
   }
 
-  private fun push(from: String, to: String, force: Boolean = false, skipHook: Boolean = false): GitPushResult {
+  private fun push(from: String, to: String, force: Boolean = false, skipHook: Boolean = false, canChangeUpstream: Boolean = false): GitPushResult {
     updateRepositories()
     refresh()
     updateChangeListManager()
 
-    val spec = makePushSpec(repository, from, to)
+    val spec = makePushSpec(repository, from, to, canChangeUpstream)
     return GitPushOperation(project, pushSupport, singletonMap(repository, spec), null, force, skipHook).execute()
   }
 

plugins/git4idea/tests/git4idea/test/GitTestUtil.kt

@@ -183,7 +183,7 @@ fun findGitLogProvider(project: Project): GitLogProvider {
   return providers[0] as GitLogProvider
 }
 
-internal fun makePushSpec(repository: GitRepository, from: String, to: String): PushSpec<GitPushSource, GitPushTarget> {
+internal fun makePushSpec(repository: GitRepository, from: String, to: String, canChangeUpstream: Boolean = false): PushSpec<GitPushSource, GitPushTarget> {
   val source = repository.branches.findLocalBranch(from)!!
   var target: GitRemoteBranch? = repository.branches.findBranchByName(to) as GitRemoteBranch?
   val newBranch: Boolean
@@ -196,7 +196,11 @@ internal fun makePushSpec(repository: GitRepository, from: String, to: String):
   else {
     newBranch = false
   }
-  return PushSpec(GitPushSource.create(source), GitPushTarget(target, newBranch))
+  val pushTarget = GitPushTarget(target, newBranch)
+  if (canChangeUpstream) {
+    pushTarget.shouldSetNewUpstream(true)
+  }
+  return PushSpec(GitPushSource.create(source), pushTarget)
 }
 
 internal fun GitRepository.resolveConflicts() {