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
12import fnmatch
13import os
14import zipfile
15
16
17def filtered(names, blacklist):
18  """Filter the list of file or directory names."""
19  rv = names[:]
20  for pattern in blacklist:
21    rv = [n for n in rv if not fnmatch.fnmatch(n, pattern)]
22  return rv
23
24
25def zip(target_dir, zip_file, blacklist=None):  # pylint: disable=W0622
26  """Zip the given directory, write to the given zip file."""
27  if not os.path.isdir(target_dir):
28    raise IOError('%s does not exist!' % target_dir)
29  blacklist = blacklist or []
30  with zipfile.ZipFile(zip_file, 'w', zipfile.ZIP_DEFLATED, True) as z:
31    for r, d, f in os.walk(target_dir, topdown=True):
32      d[:] = filtered(d, blacklist)
33      for filename in filtered(f, blacklist):
34        filepath = os.path.join(r, filename)
35        zi = zipfile.ZipInfo(filepath)
36        zi.filename = os.path.relpath(filepath, target_dir)
37        perms = os.stat(filepath).st_mode
38        zi.external_attr = perms << 16L
39        zi.compress_type = zipfile.ZIP_DEFLATED
40        with open(filepath, 'rb') as f:
41          content = f.read()
42        z.writestr(zi, content)
43      for dirname in d:
44        dirpath = os.path.join(r, dirname)
45        z.write(dirpath, os.path.relpath(dirpath, target_dir))
46
47
48def unzip(zip_file, target_dir):
49  """Unzip the given zip file into the target dir."""
50  if not os.path.isdir(target_dir):
51    os.makedirs(target_dir)
52  with zipfile.ZipFile(zip_file, 'r', zipfile.ZIP_DEFLATED, True) as z:
53    for zi in z.infolist():
54      dst_path = os.path.join(target_dir, zi.filename)
55      if zi.filename.endswith('/'):
56        os.mkdir(dst_path)
57      else:
58        with open(dst_path, 'wb') as f:
59          f.write(z.read(zi))
60      perms = zi.external_attr >> 16L
61      os.chmod(dst_path, perms)
62