Home | History | Annotate | Download | only in utils 
      1 # Copyright 2015 The Chromium Authors. All rights reserved.
      2 # Use of this source code is governed by a BSD-style license that can be
      3 # found in the LICENSE file.
      4 
      5 import argparse
      6 import json
      7 import logging
      8 import os
      9 import sys
     10 import zipfile
     11 
     12 if __name__ == '__main__':
     13   _DEVIL_ROOT_DIR = os.path.abspath(
     14       os.path.join(os.path.dirname(__file__), '..', '..'))
     15   _PY_UTILS_ROOT_DIR = os.path.abspath(
     16       os.path.join(_DEVIL_ROOT_DIR, '..', 'common', 'py_utils'))
     17   sys.path.extend((_DEVIL_ROOT_DIR, _PY_UTILS_ROOT_DIR))
     18 
     19 from devil import base_error
     20 from devil.utils import cmd_helper
     21 from py_utils import tempfile_ext
     22 
     23 
     24 logger = logging.getLogger(__name__)
     25 
     26 
     27 class ZipFailedError(base_error.BaseError):
     28   """Raised on a failure to perform a zip operation."""
     29   pass
     30 
     31 
     32 def _WriteToZipFile(zip_file, path, arc_path):
     33   """Recursively write |path| to |zip_file| as |arc_path|.
     34 
     35   zip_file: An open instance of zipfile.ZipFile.
     36   path: An absolute path to the file or directory to be zipped.
     37   arc_path: A relative path within the zip file to which the file or directory
     38     located at |path| should be written.
     39   """
     40   if os.path.isdir(path):
     41     for dir_path, _, file_names in os.walk(path):
     42       dir_arc_path = os.path.join(arc_path, os.path.relpath(dir_path, path))
     43       logger.debug('dir:  %s -> %s', dir_path, dir_arc_path)
     44       zip_file.write(dir_path, dir_arc_path, zipfile.ZIP_STORED)
     45       for f in file_names:
     46         file_path = os.path.join(dir_path, f)
     47         file_arc_path = os.path.join(dir_arc_path, f)
     48         logger.debug('file: %s -> %s', file_path, file_arc_path)
     49         zip_file.write(file_path, file_arc_path, zipfile.ZIP_DEFLATED)
     50   else:
     51     logger.debug('file: %s -> %s', path, arc_path)
     52     zip_file.write(path, arc_path, zipfile.ZIP_DEFLATED)
     53 
     54 
     55 def _WriteZipFile(zip_path, zip_contents):
     56   with zipfile.ZipFile(zip_path, 'w') as zip_file:
     57     for path, arc_path in zip_contents:
     58       _WriteToZipFile(zip_file, path, arc_path)
     59 
     60 
     61 def WriteZipFile(zip_path, zip_contents):
     62   """Writes the provided contents to the given zip file.
     63 
     64   Note that this uses python's zipfile module and is done in a separate
     65   process to avoid hogging the GIL.
     66 
     67   Args:
     68     zip_path: String path to the zip file to write.
     69     zip_contents: A list of (host path, archive path) tuples.
     70 
     71   Raises:
     72     ZipFailedError on failure.
     73   """
     74   zip_spec = {
     75     'zip_path': zip_path,
     76     'zip_contents': zip_contents,
     77   }
     78   with tempfile_ext.NamedTemporaryDirectory() as tmpdir:
     79     json_path = os.path.join(tmpdir, 'zip_spec.json')
     80     with open(json_path, 'w') as json_file:
     81       json.dump(zip_spec, json_file)
     82     ret, output, error = cmd_helper.GetCmdStatusOutputAndError([
     83         sys.executable, os.path.abspath(__file__),
     84         '--zip-spec', json_path])
     85 
     86   if ret != 0:
     87     exc_msg = ['Failed to create %s' % zip_path]
     88     exc_msg.extend('stdout:  %s' % l for l in output.splitlines())
     89     exc_msg.extend('stderr:  %s' % l for l in error.splitlines())
     90     raise ZipFailedError('\n'.join(exc_msg))
     91 
     92 
     93 def main(raw_args):
     94   parser = argparse.ArgumentParser()
     95   parser.add_argument('--zip-spec', required=True)
     96 
     97   args = parser.parse_args(raw_args)
     98 
     99   with open(args.zip_spec) as zip_spec_file:
    100     zip_spec = json.load(zip_spec_file)
    101 
    102   return _WriteZipFile(**zip_spec)
    103 
    104 
    105 if __name__ == '__main__':
    106   sys.exit(main(sys.argv[1:]))
    107