Add the ability to set the upstream branch

git_cl.py

 #!/usr/bin/env python
 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 # Copyright (C) 2008 Evan Martin <martine@danga.com>
 
 """A git-command for integrating reviews on Rietveld."""
 
 import json
@@ skipping 1496 matching lines @@
     return 2
 
   print "The tree is %s" % status
   print
   print GetTreeStatusReason()
   if status != 'open':
     return 1
   return 0
 
 
+@usage('[new upstream branch]')
 def CMDupstream(parser, args):
-  """print the name of the upstream branch, if any"""
+  """prints or sets the name of the upstream branch, if any"""
   _, args = parser.parse_args(args)
+  if len(args) > 1:
+    parser.error('Unrecognized args: %s' % ' '.join(args))
+    return 0
+
+  cl = Changelist()
   if args:
-    parser.error('Unrecognized args: %s' % ' '.join(args))
-  cl = Changelist()
-  print cl.GetUpstreamBranch()
+    # One arg means set upstream branch.
+    RunGit(['branch', '--set-upstream', cl.GetBranch(), args[0]])
+    cl = Changelist()
+    print "Upstream branch set to " + cl.GetUpstreamBranch()
+  else:
+    print cl.GetUpstreamBranch()
   return 0
 
 
 def CMDset_commit(parser, args):
   """set the commit bit"""
   _, args = parser.parse_args(args)
   if args:
     parser.error('Unrecognized args: %s' % ' '.join(args))
   cl = Changelist()
   cl.SetFlag('commit', '1')
@@ skipping 74 matching lines @@
              'and retry or visit go/isgaeup.\n%s') % (e.code, str(e)))
 
   # Not a known command. Default to help.
   GenUsage(parser, 'help')
   return CMDhelp(parser, argv)
 
 
 if __name__ == '__main__':
   fix_encoding.fix_encoding()
   sys.exit(main(sys.argv[1:]))