Add asset management scripts
These provide an easy way to create assets to be used by bots,
eg. Android SDK.
To create an asset:
$ infra/bots/assets/assets.py add android_sdk
(adds scripts in infra/bots/assets/android_sdk)
To upload a new version of an asset:
$ infra/bots/assets/android_sdk/upload.py -t $ANDROID_SDK_ROOT
(uploads Android SDK to GS, writes a version file)
$ git commit
$ git cl upload
To download the current version of the asset:
$ infra/bots/assets/android_sdk/download.py -t ../tmp
BUG=skia:5427
GOLD_TRYBOT_URL= https://gold.skia.org/search?issue=2069543002
Review-Url: https://codereview.chromium.org/2069543002
diff --git a/infra/bots/assets/README.md b/infra/bots/assets/README.md
new file mode 100644
index 0000000..91d0247
--- /dev/null
+++ b/infra/bots/assets/README.md
@@ -0,0 +1,47 @@
+Assets
+======
+
+This directory contains tooling for managing assets used by the bots. The
+primary entry point is assets.py, which allows a user to add, remove, upload,
+and download assets.
+
+Assets are stored in Google Storage, named for their version number.
+
+
+Individual Assets
+-----------------
+
+Each asset has its own subdirectory with the following contents:
+* VERSION: The current version number of the asset.
+* download.py: Convenience script for downloading the current version of the asset.
+* upload.py: Convenience script for uploading a new version of the asset.
+* [optional] create.py: Script which creates the asset, implemented by the user.
+* [optional] create\_and\_upload.py: Convenience script which combines create.py with upload.py.
+
+
+Examples
+-------
+
+Add a new asset and upload an initial version.
+
+```
+$ infra/bots/assets/assets.py add myasset
+Creating asset in infra/bots/assets/myasset
+Creating infra/bots/assets/myasset/download.py
+Creating infra/bots/assets/myasset/upload.py
+Creating infra/bots/assets/myasset/common.py
+Add script to automate creation of this asset? (y/n) n
+$ infra/bots/assets/myasset/upload.py -t ${MY_ASSET_LOCATION}
+$ git commit
+```
+
+Add an asset whose creation can be automated.
+
+```
+$ infra/bots/assets/assets.py add myasset
+Add script to automate creation of this asset? (y/n) y
+$ vi infra/bots/assets/myasset/create.py
+(implement the create_asset function)
+$ infra/bots/assets/myasset/create_and_upload.py
+$ git commit
+```
diff --git a/infra/bots/assets/__init__.py b/infra/bots/assets/__init__.py
new file mode 100644
index 0000000..78953f5
--- /dev/null
+++ b/infra/bots/assets/__init__.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
diff --git a/infra/bots/assets/asset_utils.py b/infra/bots/assets/asset_utils.py
new file mode 100644
index 0000000..a13d229
--- /dev/null
+++ b/infra/bots/assets/asset_utils.py
@@ -0,0 +1,174 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Utilities for managing assets."""
+
+
+import argparse
+import os
+import shlex
+import shutil
+import subprocess
+import sys
+
+SKIA_DIR = os.path.abspath(os.path.realpath(os.path.join(
+ os.path.dirname(os.path.abspath(__file__)),
+ os.pardir, os.pardir, os.pardir)))
+INFRA_BOTS_DIR = os.path.join(SKIA_DIR, 'infra', 'bots')
+sys.path.insert(0, INFRA_BOTS_DIR)
+import utils
+import zip_utils
+
+
+ASSETS_DIR = os.path.join(INFRA_BOTS_DIR, 'assets')
+DEFAULT_GS_BUCKET = 'skia-buildbots'
+GS_SUBDIR_TMPL = 'gs://%s/assets/%s'
+GS_PATH_TMPL = '%s/%s.zip'
+VERSION_FILENAME = 'VERSION'
+ZIP_BLACKLIST = ['.git', '.svn', '*.pyc', '.DS_STORE']
+
+
+class _GSWrapper(object):
+ """Wrapper object for interacting with Google Storage."""
+ def __init__(self, gsutil):
+ gsutil = os.path.abspath(gsutil) if gsutil else 'gsutil'
+ self._gsutil = [gsutil]
+ if gsutil.endswith('.py'):
+ self._gsutil = ['python', gsutil]
+
+ def copy(self, src, dst):
+ """Copy src to dst."""
+ subprocess.check_call(self._gsutil + ['cp', src, dst])
+
+ def list(self, path):
+ """List objects in the given path."""
+ try:
+ return subprocess.check_output(self._gsutil + ['ls', path]).splitlines()
+ except subprocess.CalledProcessError:
+ # If the prefix does not exist, we'll get an error, which is okay.
+ return []
+
+
+def _prompt(prompt):
+ """Prompt for input, return result."""
+ return raw_input(prompt)
+
+
+class Asset(object):
+ def __init__(self, name, gs_bucket=DEFAULT_GS_BUCKET, gsutil=None):
+ self._gs = _GSWrapper(gsutil)
+ self._gs_subdir = GS_SUBDIR_TMPL % (gs_bucket, name)
+ self._name = name
+ self._dir = os.path.join(ASSETS_DIR, self._name)
+
+ @property
+ def version_file(self):
+ """Return the path to the version file for this asset."""
+ return os.path.join(self._dir, VERSION_FILENAME)
+
+ def get_current_version(self):
+ """Obtain the current version of the asset."""
+ if not os.path.isfile(self.version_file):
+ return -1
+ with open(self.version_file) as f:
+ return int(f.read())
+
+ def get_available_versions(self):
+ """Return the existing version numbers for this asset."""
+ files = self._gs.list(self._gs_subdir)
+ bnames = [os.path.basename(f) for f in files]
+ suffix = '.zip'
+ versions = [int(f[:-len(suffix)]) for f in bnames if f.endswith(suffix)]
+ versions.sort()
+ return versions
+
+ def get_next_version(self):
+ """Find the next available version number for the asset."""
+ versions = self.get_available_versions()
+ if len(versions) == 0:
+ return 0
+ return versions[-1] + 1
+
+ def download_version(self, version, target_dir):
+ """Download the specified version of the asset."""
+ gs_path = GS_PATH_TMPL % (self._gs_subdir, str(version))
+ target_dir = os.path.abspath(target_dir)
+ with utils.tmp_dir():
+ zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
+ self._gs.copy(gs_path, zip_file)
+ zip_utils.unzip(zip_file, target_dir)
+
+ def download_current_version(self, target_dir):
+ """Download the version of the asset specified in its version file."""
+ v = self.get_current_version()
+ self.download_version(v, target_dir)
+
+ def upload_new_version(self, target_dir, commit=False):
+ """Upload a new version and update the version file for the asset."""
+ version = self.get_next_version()
+ target_dir = os.path.abspath(target_dir)
+ with utils.tmp_dir():
+ zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
+ zip_utils.zip(target_dir, zip_file, blacklist=ZIP_BLACKLIST)
+ gs_path = GS_PATH_TMPL % (self._gs_subdir, str(version))
+ self._gs.copy(zip_file, gs_path)
+
+ def _write_version():
+ with open(self.version_file, 'w') as f:
+ f.write(str(version))
+ subprocess.check_call([utils.GIT, 'add', self.version_file])
+
+ with utils.chdir(SKIA_DIR):
+ if commit:
+ with utils.git_branch():
+ _write_version()
+ subprocess.check_call([
+ utils.GIT, 'commit', '-m', 'Update %s version' % self._name])
+ subprocess.check_call([utils.GIT, 'cl', 'upload', '--bypass-hooks'])
+ else:
+ _write_version()
+
+ @classmethod
+ def add(cls, name, gs_bucket=DEFAULT_GS_BUCKET, gsutil=None):
+ """Add an asset."""
+ asset = cls(name, gs_bucket=gs_bucket, gsutil=gsutil)
+ if os.path.isdir(asset._dir):
+ raise Exception('Asset %s already exists!' % asset._name)
+
+ print 'Creating asset in %s' % asset._dir
+ os.mkdir(asset._dir)
+ def copy_script(script):
+ src = os.path.join(ASSETS_DIR, 'scripts', script)
+ dst = os.path.join(asset._dir, script)
+ print 'Creating %s' % dst
+ shutil.copy(src, dst)
+ subprocess.check_call([utils.GIT, 'add', dst])
+
+ for script in ('download.py', 'upload.py', 'common.py'):
+ copy_script(script)
+ resp = _prompt('Add script to automate creation of this asset? (y/n) ')
+ if resp == 'y':
+ copy_script('create.py')
+ copy_script('create_and_upload.py')
+ print 'You will need to add implementation to the creation script.'
+ print 'Successfully created asset %s.' % asset._name
+ return asset
+
+ def remove(self):
+ """Remove this asset."""
+ # Ensure that the asset exists.
+ if not os.path.isdir(self._dir):
+ raise Exception('Asset %s does not exist!' % self._name)
+
+ # Remove the asset.
+ subprocess.check_call([utils.GIT, 'rm', '-rf', self._dir])
+ if os.path.isdir(self._dir):
+ shutil.rmtree(self._dir)
+
+ # We *could* remove all uploaded versions of the asset in Google Storage but
+ # we choose not to be that destructive.
diff --git a/infra/bots/assets/asset_utils_test.py b/infra/bots/assets/asset_utils_test.py
new file mode 100644
index 0000000..edfc271
--- /dev/null
+++ b/infra/bots/assets/asset_utils_test.py
@@ -0,0 +1,124 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Tests for asset_utils."""
+
+
+import asset_utils
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import unittest
+import uuid
+
+
+FILE_DIR = os.path.dirname(os.path.abspath(__file__))
+INFRA_BOTS_DIR = os.path.realpath(os.path.join(
+ FILE_DIR, os.pardir, 'infra', 'bots'))
+sys.path.insert(0, INFRA_BOTS_DIR)
+import test_utils
+import utils
+
+
+GS_BUCKET = 'skia-infra-testdata'
+
+
+def _fake_prompt(result):
+ """Make a function that pretends to prompt for input and returns a result."""
+ return lambda s: result
+
+
+def _write_stuff(target_dir):
+ """Write some files and directories into target_dir."""
+ fw = test_utils.FileWriter(target_dir)
+ fw.mkdir('mydir')
+ fw.mkdir('anotherdir', 0666)
+ fw.mkdir('dir3', 0600)
+ fw.mkdir('subdir')
+ fw.write('a.txt', 0777)
+ fw.write('b.txt', 0751)
+ fw.write('c.txt', 0640)
+ fw.write(os.path.join('subdir', 'd.txt'), 0640)
+
+
+class AssetUtilsTest(unittest.TestCase):
+ def setUp(self):
+ self.asset_name = str(uuid.uuid4())
+ self.old_prompt = asset_utils._prompt
+ asset_utils._prompt = _fake_prompt('y')
+ self.a = asset_utils.Asset.add(self.asset_name, gs_bucket=GS_BUCKET)
+
+ def tearDown(self):
+ if self.a:
+ self.a.remove()
+ asset_utils._prompt = self.old_prompt
+
+ gs_path = 'gs://%s/assets/%s' % (GS_BUCKET, self.asset_name)
+ attempt_delete = True
+ try:
+ subprocess.check_call(['gsutil', 'ls', gs_path])
+ except subprocess.CalledProcessError:
+ attempt_delete = False
+ if attempt_delete:
+ subprocess.check_call(['gsutil', 'rm', '-rf', gs_path])
+
+ def test_add_remove(self):
+ # Ensure that we can't create an asset twice.
+ with self.assertRaises(Exception):
+ asset_utils.Asset.add(self.asset_name, gs_bucket=GS_BUCKET)
+
+ # Ensure that the asset dir exists.
+ asset_dir = os.path.join(FILE_DIR, self.asset_name)
+ self.assertTrue(os.path.isdir(asset_dir))
+
+ # Remove the asset, ensure that it's gone.
+ self.a.remove()
+ self.a = None
+ self.assertFalse(os.path.exists(asset_dir))
+
+ def test_upload_download(self):
+ with utils.tmp_dir():
+ # Create input files and directories.
+ input_dir = os.path.join(os.getcwd(), 'input')
+ _write_stuff(input_dir)
+
+ # Upload a version, download it again.
+ self.a.upload_new_version(input_dir)
+ output_dir = os.path.join(os.getcwd(), 'output')
+ self.a.download_current_version(output_dir)
+
+ # Compare.
+ test_utils.compare_trees(self, input_dir, output_dir)
+
+ def test_versions(self):
+ with utils.tmp_dir():
+ # Create input files and directories.
+ input_dir = os.path.join(os.getcwd(), 'input')
+ _write_stuff(input_dir)
+
+ self.assertEqual(self.a.get_current_version(), -1)
+ self.assertEqual(self.a.get_available_versions(), [])
+ self.assertEqual(self.a.get_next_version(), 0)
+
+ self.a.upload_new_version(input_dir)
+
+ self.assertEqual(self.a.get_current_version(), 0)
+ self.assertEqual(self.a.get_available_versions(), [0])
+ self.assertEqual(self.a.get_next_version(), 1)
+
+ self.a.upload_new_version(input_dir)
+
+ self.assertEqual(self.a.get_current_version(), 1)
+ self.assertEqual(self.a.get_available_versions(), [0, 1])
+ self.assertEqual(self.a.get_next_version(), 2)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/infra/bots/assets/assets.py b/infra/bots/assets/assets.py
new file mode 100755
index 0000000..538a41b
--- /dev/null
+++ b/infra/bots/assets/assets.py
@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Tool for managing assets."""
+
+
+import argparse
+import asset_utils
+import os
+import shutil
+import subprocess
+import sys
+
+FILE_DIR = os.path.dirname(os.path.abspath(__file__))
+INFRA_BOTS_DIR = os.path.realpath(os.path.join(FILE_DIR, os.pardir))
+
+sys.path.insert(0, INFRA_BOTS_DIR)
+import utils
+
+
+def add(args):
+ """Add a new asset."""
+ asset_utils.Asset.add(args.asset_name)
+
+
+def remove(args):
+ """Remove an asset."""
+ asset_utils.Asset(args.asset_name).remove()
+
+
+def download(args):
+ """Download the current version of an asset."""
+ asset = asset_utils.Asset(args.asset_name, gsutil=args.gsutil)
+ asset.download_current_version(args.target_dir)
+
+
+def upload(args):
+ """Upload a new version of the asset."""
+ asset = asset_utils.Asset(args.asset_name, gsutil=args.gsutil)
+ asset.upload_new_version(args.target_dir, commit=args.commit)
+
+
+def main(argv):
+ parser = argparse.ArgumentParser(description='Tool for managing assets.')
+ subs = parser.add_subparsers(help='Commands:')
+
+ prs_add = subs.add_parser('add', help='Add a new asset.')
+ prs_add.set_defaults(func=add)
+ prs_add.add_argument('asset_name', help='Name of the asset.')
+
+ prs_remove = subs.add_parser('remove', help='Remove an asset.')
+ prs_remove.set_defaults(func=remove)
+ prs_remove.add_argument('asset_name', help='Name of the asset.')
+
+ prs_download = subs.add_parser(
+ 'download', help='Download the current version of an asset.')
+ prs_download.set_defaults(func=download)
+ prs_download.add_argument('asset_name', help='Name of the asset.')
+ prs_download.add_argument('--target_dir', '-t', required=True)
+ prs_download.add_argument('--gsutil')
+
+ prs_upload = subs.add_parser(
+ 'upload', help='Upload a new version of an asset.')
+ prs_upload.set_defaults(func=upload)
+ prs_upload.add_argument('asset_name', help='Name of the asset.')
+ prs_upload.add_argument('--target_dir', '-t', required=True)
+ prs_upload.add_argument('--gsutil')
+ prs_upload.add_argument('--commit', action='store_true')
+
+ args = parser.parse_args(argv)
+ args.func(args)
+
+
+if __name__ == '__main__':
+ main(sys.argv[1:])
diff --git a/infra/bots/assets/scripts/common.py b/infra/bots/assets/scripts/common.py
new file mode 100755
index 0000000..4920c9b
--- /dev/null
+++ b/infra/bots/assets/scripts/common.py
@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Common vars used by scripts in this directory."""
+
+
+import os
+import sys
+
+FILE_DIR = os.path.dirname(os.path.abspath(__file__))
+INFRA_BOTS_DIR = os.path.realpath(os.path.join(FILE_DIR, os.pardir, os.pardir))
+
+sys.path.insert(0, INFRA_BOTS_DIR)
+from assets import assets
+
+ASSET_NAME = os.path.basename(FILE_DIR)
+
+
+def run(cmd):
+ """Run a command, eg. "upload" or "download". """
+ assets.main([cmd, ASSET_NAME] + sys.argv[1:])
diff --git a/infra/bots/assets/scripts/create.py b/infra/bots/assets/scripts/create.py
new file mode 100755
index 0000000..4f17608
--- /dev/null
+++ b/infra/bots/assets/scripts/create.py
@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Create the asset."""
+
+
+import argparse
+
+
+def create_asset(target_dir):
+ """Create the asset."""
+ raise NotImplementedError('Implement me!')
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--target_dir', '-t', required=True)
+ args = parser.parse_args()
+ create_asset(args.target_dir)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/infra/bots/assets/scripts/create_and_upload.py b/infra/bots/assets/scripts/create_and_upload.py
new file mode 100755
index 0000000..1356447
--- /dev/null
+++ b/infra/bots/assets/scripts/create_and_upload.py
@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Create the asset and upload it."""
+
+
+import argparse
+import common
+import os
+import subprocess
+import sys
+import utils
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--gsutil')
+ args = parser.parse_args()
+
+ with utils.tmp_dir():
+ cwd = os.getcwd()
+ create_script = os.path.join(common.FILE_DIR, 'create.py')
+ upload_script = os.path.join(common.FILE_DIR, 'upload.py')
+
+ try:
+ subprocess.check_call(['python', create_script, '-t', cwd])
+ cmd = ['python', upload_script, '-t', cwd]
+ if args.gsutil:
+ cmd.extend(['--gsutil', args.gsutil])
+ subprocess.check_call(cmd)
+ except subprocess.CalledProcessError:
+ # Trap exceptions to avoid printing two stacktraces.
+ sys.exit(1)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/infra/bots/assets/scripts/download.py b/infra/bots/assets/scripts/download.py
new file mode 100755
index 0000000..96cc87d
--- /dev/null
+++ b/infra/bots/assets/scripts/download.py
@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Download the current version of the asset."""
+
+
+import common
+
+
+if __name__ == '__main__':
+ common.run('download')
diff --git a/infra/bots/assets/scripts/upload.py b/infra/bots/assets/scripts/upload.py
new file mode 100755
index 0000000..ba7fc8b
--- /dev/null
+++ b/infra/bots/assets/scripts/upload.py
@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Upload a new version of the asset."""
+
+
+import common
+
+
+if __name__ == '__main__':
+ common.run('upload')
diff --git a/infra/bots/download_asset.isolate b/infra/bots/download_asset.isolate
new file mode 100644
index 0000000..0c4d873
--- /dev/null
+++ b/infra/bots/download_asset.isolate
@@ -0,0 +1,10 @@
+{
+ 'includes': [
+ 'infrabots.isolate',
+ ],
+ 'variables': {
+ 'command': [
+ 'python', 'assets/<(ASSET)/download.py', '-t', '${ISOLATED_OUTDIR}', '--gsutil', '<(GSUTIL)',
+ ],
+ },
+}
diff --git a/infra/bots/test_utils.py b/infra/bots/test_utils.py
new file mode 100644
index 0000000..aa7ac0f
--- /dev/null
+++ b/infra/bots/test_utils.py
@@ -0,0 +1,73 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Test utilities."""
+
+
+import filecmp
+import os
+import uuid
+
+
+class FileWriter(object):
+ """Write files into a given directory."""
+ def __init__(self, cwd):
+ self._cwd = cwd
+ if not os.path.exists(self._cwd):
+ os.makedirs(self._cwd)
+
+ def mkdir(self, dname, mode=0755):
+ """Create the given directory with the given mode."""
+ dname = os.path.join(self._cwd, dname)
+ os.mkdir(dname)
+ os.chmod(dname, mode)
+
+ def write(self, fname, mode=0640):
+ """Write the file with the given mode and random contents."""
+ fname = os.path.join(self._cwd, fname)
+ with open(fname, 'w') as f:
+ f.write(str(uuid.uuid4()))
+ os.chmod(fname, mode)
+
+ def remove(self, fname):
+ """Remove the file."""
+ fname = os.path.join(self._cwd, fname)
+ if os.path.isfile(fname):
+ os.remove(fname)
+ else:
+ os.rmdir(fname)
+
+
+def compare_trees(test, a, b):
+ """Compare two directory trees, assert if any differences."""
+ def _cmp(prefix, dcmp):
+ # Verify that the file and directory names are the same.
+ test.assertEqual(len(dcmp.left_only), 0)
+ test.assertEqual(len(dcmp.right_only), 0)
+ test.assertEqual(len(dcmp.diff_files), 0)
+ test.assertEqual(len(dcmp.funny_files), 0)
+
+ # Verify that the files are identical.
+ for f in dcmp.common_files:
+ pathA = os.path.join(a, prefix, f)
+ pathB = os.path.join(b, prefix, f)
+ test.assertTrue(filecmp.cmp(pathA, pathB, shallow=False))
+ statA = os.stat(pathA)
+ statB = os.stat(pathB)
+ test.assertEqual(statA.st_mode, statB.st_mode)
+ with open(pathA, 'rb') as f:
+ contentsA = f.read()
+ with open(pathB, 'rb') as f:
+ contentsB = f.read()
+ test.assertEqual(contentsA, contentsB)
+
+ # Recurse on subdirectories.
+ for prefix, obj in dcmp.subdirs.iteritems():
+ _cmp(prefix, obj)
+
+ _cmp('', filecmp.dircmp(a, b))
diff --git a/infra/bots/zip_utils.py b/infra/bots/zip_utils.py
new file mode 100644
index 0000000..7f269b9
--- /dev/null
+++ b/infra/bots/zip_utils.py
@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Utilities for zipping and unzipping files."""
+
+
+import fnmatch
+import os
+import zipfile
+
+
+def filtered(names, blacklist):
+ """Filter the list of file or directory names."""
+ rv = names[:]
+ for pattern in blacklist:
+ rv = [n for n in rv if not fnmatch.fnmatch(n, pattern)]
+ return rv
+
+
+def zip(target_dir, zip_file, blacklist=None): # pylint: disable=W0622
+ """Zip the given directory, write to the given zip file."""
+ if not os.path.isdir(target_dir):
+ raise IOError('%s does not exist!' % target_dir)
+ blacklist = blacklist or []
+ with zipfile.ZipFile(zip_file, 'w') as z:
+ for r, d, f in os.walk(target_dir, topdown=True):
+ d[:] = filtered(d, blacklist)
+ for filename in filtered(f, blacklist):
+ filepath = os.path.join(r, filename)
+ zi = zipfile.ZipInfo(filepath)
+ zi.filename = os.path.relpath(filepath, target_dir)
+ perms = os.stat(filepath).st_mode
+ zi.external_attr = perms << 16L
+ zi.compress_type = zipfile.ZIP_STORED
+ with open(filepath, 'rb') as f:
+ content = f.read()
+ z.writestr(zi, content)
+ for dirname in d:
+ dirpath = os.path.join(r, dirname)
+ z.write(dirpath, os.path.relpath(dirpath, target_dir))
+
+
+def unzip(zip_file, target_dir):
+ """Unzip the given zip file into the target dir."""
+ if not os.path.isdir(target_dir):
+ os.makedirs(target_dir)
+ with zipfile.ZipFile(zip_file, 'r') as z:
+ for zi in z.infolist():
+ dst_path = os.path.join(target_dir, zi.filename)
+ if zi.filename.endswith('/'):
+ os.mkdir(dst_path)
+ else:
+ with open(dst_path, 'w') as f:
+ f.write(z.read(zi))
+ perms = zi.external_attr >> 16L
+ os.chmod(dst_path, perms)
diff --git a/infra/bots/zip_utils_test.py b/infra/bots/zip_utils_test.py
new file mode 100644
index 0000000..4f88a11
--- /dev/null
+++ b/infra/bots/zip_utils_test.py
@@ -0,0 +1,74 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Tests for zip_utils."""
+
+
+import filecmp
+import os
+import test_utils
+import unittest
+import utils
+import uuid
+import zip_utils
+
+
+class ZipUtilsTest(unittest.TestCase):
+ def test_zip_unzip(self):
+ with utils.tmp_dir():
+ fw = test_utils.FileWriter(os.path.join(os.getcwd(), 'input'))
+ # Create input files and directories.
+ fw.mkdir('mydir')
+ fw.mkdir('anotherdir', 0666)
+ fw.mkdir('dir3', 0600)
+ fw.mkdir('subdir')
+ fw.write('a.txt', 0777)
+ fw.write('b.txt', 0751)
+ fw.write('c.txt', 0640)
+ fw.write(os.path.join('subdir', 'd.txt'), 0640)
+
+ # Zip, unzip.
+ zip_utils.zip('input', 'test.zip')
+ zip_utils.unzip('test.zip', 'output')
+
+ # Compare the inputs and outputs.
+ test_utils.compare_trees(self, 'input', 'output')
+
+ def test_blacklist(self):
+ with utils.tmp_dir():
+ # Create input files and directories.
+ fw = test_utils.FileWriter(os.path.join(os.getcwd(), 'input'))
+ fw.mkdir('.git')
+ fw.write(os.path.join('.git', 'index'))
+ fw.write('somefile')
+ fw.write('.DS_STORE')
+ fw.write('leftover.pyc')
+ fw.write('.pycfile')
+
+ # Zip, unzip.
+ zip_utils.zip('input', 'test.zip', blacklist=['.git', '.DS*', '*.pyc'])
+ zip_utils.unzip('test.zip', 'output')
+
+ # Remove the files/dirs we don't expect to see in output, so that we can
+ # use self._compare_trees to check the results.
+ fw.remove(os.path.join('.git', 'index'))
+ fw.remove('.git')
+ fw.remove('.DS_STORE')
+ fw.remove('leftover.pyc')
+
+ # Compare results.
+ test_utils.compare_trees(self, 'input', 'output')
+
+ def test_nonexistent_dir(self):
+ with utils.tmp_dir():
+ with self.assertRaises(IOError):
+ zip_utils.zip('input', 'test.zip')
+
+
+if __name__ == '__main__':
+ unittest.main()