Handle KeyboardInterrupt consistently in python scripts

git_map.py

 #!/usr/bin/env python
 # Copyright 2014 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.
 
 """
 Provides an augmented `git log --graph` view. In particular, it also annotates
 commits with branches + tags that point to them. Items are colorized as follows:
   * Cyan    - Currently checked out branch
   * Green   - Local branch
@@ skipping 19 matching lines @@
 WHITE = colorama.Fore.WHITE
 
 BLUEBAK = colorama.Back.BLUE
 
 BRIGHT = colorama.Style.BRIGHT
 RESET = colorama.Fore.RESET + colorama.Back.RESET + colorama.Style.RESET_ALL
 
 # Git emits combined color
 BRIGHT_RED = '\x1b[1;31m'
 
-def main():
+def main(argv):
   map_extra = config_list('depot_tools.map_extra')
   fmt = '%C(red bold)%h%x09%Creset%C(green)%d%Creset %C(yellow)%ad%Creset ~ %s'
   log_proc = subprocess2.Popen(
     [GIT_EXE, 'log', '--graph', '--branches', '--tags', root(),
      '--color=always', '--date=short', ('--pretty=format:' + fmt)
-    ] + map_extra + sys.argv[1:],
+    ] + map_extra + argv,
     stdout=subprocess2.PIPE,
     shell=False)
 
   current = current_branch()
   all_branches = set(branches())
   merge_base_map = {b: get_or_create_merge_base(b) for b in all_branches}
   merge_base_map = {b: v for b, v in merge_base_map.iteritems() if v}
   if current in all_branches:
     all_branches.remove(current)
   all_tags = set(tags())
@@ skipping 46 matching lines @@
       sys.stdout.write(line)
   except (IOError, KeyboardInterrupt):
     pass
   finally:
     sys.stderr.close()
     sys.stdout.close()
   return 0
 
 
 if __name__ == '__main__':
-  sys.exit(main())
-
+  try:
+    sys.exit(main(sys.argv[1:]))
+  except KeyboardInterrupt:
+    sys.stderr.write('interrupted\n')
+    sys.exit(1)