Merge pull request #559 from josephfrazier/git-push-explicit-upstream

Fix suggestions for `git push -u origin`

tests/rules/test_git_push.py

@@ -23,5 +23,9 @@ def test_match(stderr):
 def test_get_new_command(stderr):
     assert get_new_command(Command('git push', stderr=stderr))\
         == "git push --set-upstream origin master"
+    assert get_new_command(Command('git push -u origin', stderr=stderr))\
+        == "git push --set-upstream origin master"
+    assert get_new_command(Command('git push --set-upstream origin', stderr=stderr))\
+        == "git push --set-upstream origin master"
     assert get_new_command(Command('git push --quiet', stderr=stderr))\
         == "git push --set-upstream origin master --quiet"

thefuck/rules/git_push.py

@@ -10,5 +10,21 @@ def match(command):
 
 @git_support
 def get_new_command(command):
+    # If --set-upstream or -u are passed, remove it and its argument. This is
+    # because the remaining arguments are concatenated onto the command suggested
+    # by git, which includes --set-upstream and its argument
+    upstream_option_index = -1
+    try:
+        upstream_option_index = command.script_parts.index('--set-upstream')
+    except ValueError:
+        pass
+    try:
+        upstream_option_index = command.script_parts.index('-u')
+    except ValueError:
+        pass
+    if upstream_option_index is not -1:
+        command.script_parts.pop(upstream_option_index)
+        command.script_parts.pop(upstream_option_index)
+
     push_upstream = command.stderr.split('\n')[-3].strip().partition('git ')[2]
-    return replace_argument(command.script, 'push', push_upstream)
+    return replace_argument(" ".join(command.script_parts), 'push', push_upstream)