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 ntpath
14import os
15import posixpath
16import zipfile
17
18
19def 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
27def 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 []
32  with zipfile.ZipFile(zip_file, 'w', zipfile.ZIP_DEFLATED, True) as z:
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)
39        if os.name == 'nt':
40          # Dumb path separator replacement for Windows.
41          zi.filename = zi.filename.replace(ntpath.sep, posixpath.sep)
42        perms = os.stat(filepath).st_mode
43        zi.external_attr = perms << 16L
44        zi.compress_type = zipfile.ZIP_DEFLATED
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
53def 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)
57  with zipfile.ZipFile(zip_file, 'r', zipfile.ZIP_DEFLATED, True) as z:
58    for zi in z.infolist():
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):
65        os.mkdir(dst_path)
66      else:
67        with open(dst_path, 'wb') as f:
68          f.write(z.read(zi))
69      perms = zi.external_attr >> 16L
70      os.chmod(dst_path, perms)
71