Adds --auto_platform to download_from_google_storage.

download_from_google_storage.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.
 
 """Download files from Google Storage based on SHA1 sums."""
 
 
 import hashlib
 import optparse
 import os
 import Queue
 import re
 import stat
 import sys
 import threading
 import time
 
 import subprocess2
 
 
 GSUTIL_DEFAULT_PATH = os.path.join(
     os.path.dirname(os.path.abspath(__file__)),
     'third_party', 'gsutil', 'gsutil')
+# Maps sys.platform to what we actually want to call them.
+PLATFORM_MAPPING = {
+    'cygwin': 'win',
+    'darwin': 'mac',
+    'linux2': 'linux',
+    'win32': 'win',
+}
 
 
 class FileNotFoundError(IOError):
   pass
 
 
 class InvalidFileError(IOError):
   pass
 
 
+class InvalidPlatformError(Exception):
+  pass
+
+
 def GetNormalizedPlatform():
   """Returns the result of sys.platform accounting for cygwin.
   Under cygwin, this will always return "win32" like the native Python."""
   if sys.platform == 'cygwin':
     return 'win32'
   return sys.platform
 
 
 # Common utilities
 class Gsutil(object):
@@ skipping 64 matching lines @@
   if code != 0:
     print >> sys.stderr, ls_err
   if code == 403:
     print >> sys.stderr, 'Got error 403 while authenticating to %s.' % base_url
     print >> sys.stderr, 'Try running "download_from_google_storage --config".'
   elif code == 404:
     print >> sys.stderr, '%s not found.' % base_url
   return (base_url, code)
 
 
+def check_platform(target):
+  """Checks if any parent directory of target matches (win|mac|linux)."""
+  assert os.path.isabs(target)
+  root, target_name = os.path.split(target)
+  if not target_name:
+    return None
+  if target_name in ('linux', 'mac', 'win'):
+    return target_name
+  return check_platform(root)
+
+
 def get_sha1(filename):
   sha1 = hashlib.sha1()
   with open(filename, 'rb') as f:
     while True:
       # Read in 1mb chunks, so it doesn't all have to be loaded into memory.
       chunk = f.read(1024*1024)
       if not chunk:
         break
       sha1.update(chunk)
   return sha1.hexdigest()
 
 
 # Download-specific code starts here
 
 def enumerate_work_queue(input_filename, work_queue, directory,
-                         recursive, ignore_errors, output, sha1_file):
+                         recursive, ignore_errors, output, sha1_file,
+                         auto_platform):
   if sha1_file:
     if not os.path.exists(input_filename):
       if not ignore_errors:
         raise FileNotFoundError('%s not found.' % input_filename)
       print >> sys.stderr, '%s not found.' % input_filename
     with open(input_filename, 'rb') as f:
       sha1_match = re.match('^([A-Za-z0-9]{40})$', f.read(1024).rstrip())
       if sha1_match:
         work_queue.put(
             (sha1_match.groups(1)[0], input_filename.replace('.sha1', '')))
@@ skipping 12 matching lines @@
     if not recursive:
       for item in dirs[:]:
         dirs.remove(item)
     else:
       for exclude in ['.svn', '.git']:
         if exclude in dirs:
           dirs.remove(exclude)
     for filename in files:
       full_path = os.path.join(root, filename)
       if full_path.endswith('.sha1'):
+        if auto_platform:
+          # Skip if the platform does not match.
+          target_platform = check_platform(os.path.abspath(full_path))
+          if not target_platform:
+            err = ('--auto_platform passed in but no platform name found in '
+                   'the path of %s' % full_path)
+            if not ignore_errors:
+              raise InvalidFileError(err)
+            print >> sys.stderr, err
+            continue
+          current_platform = PLATFORM_MAPPING[sys.platform]
+          if current_platform != target_platform:
+            continue
+
         with open(full_path, 'rb') as f:
           sha1_match = re.match('^([A-Za-z0-9]{40})$', f.read(1024).rstrip())
         if sha1_match:
           work_queue.put(
               (sha1_match.groups(1)[0], full_path.replace('.sha1', '')))
           work_queue_size += 1
         else:
           if not ignore_errors:
             raise InvalidFileError('No sha1 sum found in %s.' % filename)
           print >> sys.stderr, 'No sha1 sum found in %s.' % filename
@@ skipping 56 matching lines @@
   while True:
     line = output_queue.get()
     # Its plausible we want to print empty lines.
     if line is None:
       break
     print line
 
 
 def download_from_google_storage(
     input_filename, base_url, gsutil, num_threads, directory, recursive,
-    force, output, ignore_errors, sha1_file, verbose):
+    force, output, ignore_errors, sha1_file, verbose, auto_platform):
   # Start up all the worker threads.
   all_threads = []
   download_start = time.time()
   stdout_queue = Queue.Queue()
   work_queue = Queue.Queue()
   ret_codes = Queue.Queue()
   ret_codes.put((0, None))
   for thread_num in range(num_threads):
     t = threading.Thread(
         target=_downloader_worker_thread,
         args=[thread_num, work_queue, force, base_url,
               gsutil, stdout_queue, ret_codes, verbose])
     t.daemon = True
     t.start()
     all_threads.append(t)
   printer_thread = threading.Thread(target=printer_worker, args=[stdout_queue])
   printer_thread.daemon = True
   printer_thread.start()
 
   # Enumerate our work queue.
   work_queue_size = enumerate_work_queue(
       input_filename, work_queue, directory, recursive,
-      ignore_errors, output, sha1_file)
+      ignore_errors, output, sha1_file, auto_platform)
   for _ in all_threads:
     work_queue.put((None, None))  # Used to tell worker threads to stop.
 
   # Wait for all downloads to finish.
   for t in all_threads:
     t.join()
   stdout_queue.put(None)
   printer_thread.join()
 
   # See if we ran into any errors.
@@ skipping 50 matching lines @@
                          'to initialize your saved Google Storage '
                          'credentials.  This will create a read-only '
                          'credentials file in ~/.boto.depot_tools.')
   parser.add_option('-n', '--no_auth', action='store_true',
                     help='Skip auth checking.  Use if it\'s known that the '
                          'target bucket is a public bucket.')
   parser.add_option('-p', '--platform',
                     help='A regular expression that is compared against '
                          'Python\'s sys.platform. If this option is specified, '
                          'the download will happen only if there is a match.')
+  parser.add_option('-a', '--auto_platform',
+                    action='store_true',
+                    help='Detects if any parent folder of the target matches '
+                         '(linux|mac|win).  If so, the script will only '
+                         'process files that are in the paths that '
+                         'that matches the current platform.')
   parser.add_option('-v', '--verbose', action='store_true',
                     help='Output extra diagnostic and progress information.')
 
   (options, args) = parser.parse_args()
 
   # Make sure we should run at all based on platform matching.
   if options.platform:
+    if options.auto_platform:
+      parser.error('--platform can not be specified with --auto_platform')
     if not re.match(options.platform, GetNormalizedPlatform()):
       if options.verbose:
         print('The current platform doesn\'t match "%s", skipping.' %
               options.platform)
       return 0
 
   # Set the boto file to /dev/null if we don't need auth.
   if options.no_auth:
     options.boto = os.devnull
 
@@ skipping 17 matching lines @@
     parser.error('Too many targets.')
   if not options.bucket:
     parser.error('Missing bucket.  Specify bucket with --bucket.')
   if options.sha1_file and options.directory:
     parser.error('Both --directory and --sha1_file are specified, '
                  'can only specify one.')
   if options.recursive and not options.directory:
     parser.error('--recursive specified but --directory not specified.')
   if options.output and options.directory:
     parser.error('--directory is specified, so --output has no effect.')
+  if (not (options.sha1_file or options.directory)
+      and options.auto_platform):
+    parser.error('--auto_platform must be specified with either '
+                 '--sha1_file or --directory')
+
   input_filename = args[0]
 
   # Set output filename if not specified.
   if not options.output and not options.directory:
     if not options.sha1_file:
       # Target is a sha1 sum, so output filename would also be the sha1 sum.
       options.output = input_filename
     elif options.sha1_file:
       # Target is a .sha1 file.
       if not input_filename.endswith('.sha1'):
@@ skipping 12 matching lines @@
                    % options.output)
 
   # Check we have a valid bucket with valid permissions.
   base_url, code = check_bucket_permissions(options.bucket, gsutil)
   if code:
     return code
 
   return download_from_google_storage(
       input_filename, base_url, gsutil, options.num_threads, options.directory,
       options.recursive, options.force, options.output, options.ignore_errors,
-      options.sha1_file, options.verbose)
+      options.sha1_file, options.verbose, options.auto_platform)
 
 
 if __name__ == '__main__':
   sys.exit(main(sys.argv))