diff options
Diffstat (limited to 'catapult/devil/devil/utils/zip_utils.py')
-rw-r--r-- | catapult/devil/devil/utils/zip_utils.py | 75 |
1 files changed, 74 insertions, 1 deletions
diff --git a/catapult/devil/devil/utils/zip_utils.py b/catapult/devil/devil/utils/zip_utils.py index eaa6a2df..e1f812b7 100644 --- a/catapult/devil/devil/utils/zip_utils.py +++ b/catapult/devil/devil/utils/zip_utils.py @@ -2,14 +2,34 @@ # Use of this source code is governed by a BSD-style license that can be # found in the LICENSE file. +import argparse +import json import logging import os +import sys import zipfile +if __name__ == '__main__': + _DEVIL_ROOT_DIR = os.path.abspath( + os.path.join(os.path.dirname(__file__), '..', '..')) + _PY_UTILS_ROOT_DIR = os.path.abspath( + os.path.join(_DEVIL_ROOT_DIR, '..', 'common', 'py_utils')) + sys.path.extend((_DEVIL_ROOT_DIR, _PY_UTILS_ROOT_DIR)) + +from devil import base_error +from devil.utils import cmd_helper +from py_utils import tempfile_ext + + logger = logging.getLogger(__name__) -def WriteToZipFile(zip_file, path, arc_path): +class ZipFailedError(base_error.BaseError): + """Raised on a failure to perform a zip operation.""" + pass + + +def _WriteToZipFile(zip_file, path, arc_path): """Recursively write |path| to |zip_file| as |arc_path|. zip_file: An open instance of zipfile.ZipFile. @@ -31,3 +51,56 @@ def WriteToZipFile(zip_file, path, arc_path): logger.debug('file: %s -> %s', path, arc_path) zip_file.write(path, arc_path, zipfile.ZIP_DEFLATED) + +def _WriteZipFile(zip_path, zip_contents): + with zipfile.ZipFile(zip_path, 'w') as zip_file: + for path, arc_path in zip_contents: + _WriteToZipFile(zip_file, path, arc_path) + + +def WriteZipFile(zip_path, zip_contents): + """Writes the provided contents to the given zip file. + + Note that this uses python's zipfile module and is done in a separate + process to avoid hogging the GIL. + + Args: + zip_path: String path to the zip file to write. + zip_contents: A list of (host path, archive path) tuples. + + Raises: + ZipFailedError on failure. + """ + zip_spec = { + 'zip_path': zip_path, + 'zip_contents': zip_contents, + } + with tempfile_ext.NamedTemporaryDirectory() as tmpdir: + json_path = os.path.join(tmpdir, 'zip_spec.json') + with open(json_path, 'w') as json_file: + json.dump(zip_spec, json_file) + ret, output, error = cmd_helper.GetCmdStatusOutputAndError([ + sys.executable, os.path.abspath(__file__), + '--zip-spec', json_path]) + + if ret != 0: + exc_msg = ['Failed to create %s' % zip_path] + exc_msg.extend('stdout: %s' % l for l in output.splitlines()) + exc_msg.extend('stderr: %s' % l for l in error.splitlines()) + raise ZipFailedError('\n'.join(exc_msg)) + + +def main(raw_args): + parser = argparse.ArgumentParser() + parser.add_argument('--zip-spec', required=True) + + args = parser.parse_args(raw_args) + + with open(args.zip_spec) as zip_spec_file: + zip_spec = json.load(zip_spec_file) + + return _WriteZipFile(**zip_spec) + + +if __name__ == '__main__': + sys.exit(main(sys.argv[1:])) |