borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 1 | #!/usr/bin/env python |
| 2 | # |
| 3 | # Copyright 2016 Google Inc. |
| 4 | # |
| 5 | # Use of this source code is governed by a BSD-style license that can be |
| 6 | # found in the LICENSE file. |
| 7 | |
| 8 | |
| 9 | """Utilities for zipping and unzipping files.""" |
| 10 | |
| 11 | |
| 12 | import fnmatch |
Eric Boren | 0979006 | 2017-09-28 12:31:32 -0400 | [diff] [blame] | 13 | import ntpath |
borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 14 | import os |
Eric Boren | 0979006 | 2017-09-28 12:31:32 -0400 | [diff] [blame] | 15 | import posixpath |
borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 16 | import zipfile |
| 17 | |
| 18 | |
| 19 | def filtered(names, blacklist): |
| 20 | """Filter the list of file or directory names.""" |
| 21 | rv = names[:] |
| 22 | for pattern in blacklist: |
| 23 | rv = [n for n in rv if not fnmatch.fnmatch(n, pattern)] |
| 24 | return rv |
| 25 | |
| 26 | |
| 27 | def zip(target_dir, zip_file, blacklist=None): # pylint: disable=W0622 |
| 28 | """Zip the given directory, write to the given zip file.""" |
| 29 | if not os.path.isdir(target_dir): |
| 30 | raise IOError('%s does not exist!' % target_dir) |
| 31 | blacklist = blacklist or [] |
mtklein | 5580c69 | 2016-08-26 10:52:19 -0700 | [diff] [blame] | 32 | with zipfile.ZipFile(zip_file, 'w', zipfile.ZIP_DEFLATED, True) as z: |
borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 33 | for r, d, f in os.walk(target_dir, topdown=True): |
| 34 | d[:] = filtered(d, blacklist) |
| 35 | for filename in filtered(f, blacklist): |
| 36 | filepath = os.path.join(r, filename) |
| 37 | zi = zipfile.ZipInfo(filepath) |
| 38 | zi.filename = os.path.relpath(filepath, target_dir) |
Eric Boren | 0979006 | 2017-09-28 12:31:32 -0400 | [diff] [blame] | 39 | if os.name == 'nt': |
| 40 | # Dumb path separator replacement for Windows. |
| 41 | zi.filename = zi.filename.replace(ntpath.sep, posixpath.sep) |
borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 42 | perms = os.stat(filepath).st_mode |
| 43 | zi.external_attr = perms << 16L |
mtklein | 5580c69 | 2016-08-26 10:52:19 -0700 | [diff] [blame] | 44 | zi.compress_type = zipfile.ZIP_DEFLATED |
borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 45 | with open(filepath, 'rb') as f: |
| 46 | content = f.read() |
| 47 | z.writestr(zi, content) |
| 48 | for dirname in d: |
| 49 | dirpath = os.path.join(r, dirname) |
| 50 | z.write(dirpath, os.path.relpath(dirpath, target_dir)) |
| 51 | |
| 52 | |
| 53 | def unzip(zip_file, target_dir): |
| 54 | """Unzip the given zip file into the target dir.""" |
| 55 | if not os.path.isdir(target_dir): |
| 56 | os.makedirs(target_dir) |
mtklein | 5580c69 | 2016-08-26 10:52:19 -0700 | [diff] [blame] | 57 | with zipfile.ZipFile(zip_file, 'r', zipfile.ZIP_DEFLATED, True) as z: |
borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 58 | for zi in z.infolist(): |
Eric Boren | 0979006 | 2017-09-28 12:31:32 -0400 | [diff] [blame] | 59 | dst_subpath = zi.filename |
| 60 | if os.name == 'nt': |
| 61 | # Dumb path separator replacement for Windows. |
| 62 | dst_subpath = dst_subpath.replace(posixpath.sep, ntpath.sep) |
| 63 | dst_path = os.path.join(target_dir, dst_subpath) |
| 64 | if dst_path.endswith(os.path.sep): |
borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 65 | os.mkdir(dst_path) |
| 66 | else: |
Brian Osman | 18b8490 | 2016-10-31 15:28:10 -0400 | [diff] [blame] | 67 | with open(dst_path, 'wb') as f: |
borenet | 0f1469b | 2016-06-15 12:07:42 -0700 | [diff] [blame] | 68 | f.write(z.read(zi)) |
| 69 | perms = zi.external_attr >> 16L |
| 70 | os.chmod(dst_path, perms) |