Add tree artifact support to pkg_zip

pkg/private/zip/build_zip.py

@@ -16,6 +16,7 @@
 import argparse
 import datetime
 import json
+import os
 import zipfile
 
 from pkg.private import build_info
@@ -68,53 +69,144 @@ def parse_date(ts):
   ts = datetime.datetime.utcfromtimestamp(ts)
   return (ts.year, ts.month, ts.day, ts.hour, ts.minute, ts.second)
 
-def _add_manifest_entry(options, zip_file, entry, default_mode, ts):
-  """Add an entry to the zip file.
-
-  Args:
-    options: parsed options
-    zip_file: ZipFile to write to
-    entry: manifest entry
-    default_mode: (int) file mode to use if not specified in the entry.
-    ts: (int) time stamp to add to files
-  """
-
-  entry_type, dest, src, mode, user, group = entry
-
-  # Use the pkg_tar mode/owner remaping as a fallback
-  non_abs_path = dest.strip('/')
-  dst_path = _combine_paths(options.directory, non_abs_path)
-  if entry_type == manifest.ENTRY_IS_DIR and not dst_path.endswith('/'):
-    dst_path += '/'
-  entry_info = zipfile.ZipInfo(filename=dst_path, date_time=ts)
-  # See http://www.pkware.com/documents/casestudies/APPNOTE.TXT
-  # denotes UTF-8 encoded file name.
-  entry_info.flag_bits |= 0x800
-  if mode:
-    f_mode = int(mode, 8)
-  else:
-    f_mode = default_mode
-
-  # See: https://trac.edgewall.org/attachment/ticket/8919/ZipDownload.patch
-  # external_attr is 4 bytes in size. The high order two bytes represent UNIX
-  # permission and file type bits, while the low order two contain MS-DOS FAT file
-  # attributes.
-  entry_info.external_attr = f_mode << 16
-  if entry_type == manifest.ENTRY_IS_FILE:
-    entry_info.compress_type = zipfile.ZIP_DEFLATED
-    with open(src, 'rb') as src:
-      zip_file.writestr(entry_info, src.read())
-  elif entry_type == manifest.ENTRY_IS_DIR:
-    entry_info.compress_type = zipfile.ZIP_STORED
-    # Set directory bits
-    entry_info.external_attr |= (UNIX_DIR_BIT << 16) | MSDOS_DIR_BIT
-    zip_file.writestr(entry_info, '')
-  elif entry_type == manifest.ENTRY_IS_LINK:
-    entry_info.compress_type = zipfile.ZIP_STORED
-    # Set directory bits
-    entry_info.external_attr |= (UNIX_SYMLINK_BIT << 16)
-    zip_file.writestr(entry_info, src)
-  # TODO(#309): All the rest
+
+class ZipWriter(object):
+
+  def __init__(self, output_path: str, time_stamp: int, default_mode: int):
+    """Create a writer.
+
+    You must close() after use or use in a 'with' statement.
+
+    Args:
+      output_path: path to write to
+      time_stamp: time stamp to add to files
+      default_mode: file mode to use if not specified in the entry.
+    """
+    self.output_path = output_path
+    self.time_stamp = time_stamp
+    self.default_mode = default_mode
+    self.zip_file = zipfile.ZipFile(self.output_path, mode='w')
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, t, v, traceback):
+    self.close()
+
+  def close(self):
+    self.zip_file.close()
+    self.zip_file = None
+
+  def make_zipinfo(self, path: str, mode: int):
+    """Create a Zipinfo.
+
+    Args:
+      path: file path
+      mode: file mode
+    """
+    entry_info = zipfile.ZipInfo(filename=path, date_time=self.time_stamp)
+    # See http://www.pkware.com/documents/casestudies/APPNOTE.TXT
+    # denotes UTF-8 encoded file name.
+    entry_info.flag_bits |= 0x800
+
+    # See: https://trac.edgewall.org/attachment/ticket/8919/ZipDownload.patch
+    # external_attr is 4 bytes in size. The high order two bytes represent UNIX
+    # permission and file type bits, while the low order two contain MS-DOS FAT
+    # file attributes.
+    if mode:
+      f_mode = int(mode, 8)
+    else:
+      f_mode = self.default_mode
+    entry_info.external_attr = f_mode << 16
+    return entry_info
+
+  def add_manifest_entry(self, options, entry):
+    """Add an entry to the zip file.
+
+    Args:
+      options: parsed options
+      zip_file: ZipFile to write to
+      entry: manifest entry
+    """
+
+    entry_type, dest, src, mode, user, group = entry
+
+    # Use the pkg_tar mode/owner remaping as a fallback
+    non_abs_path = dest.strip('/')
+    dst_path = _combine_paths(options.directory, non_abs_path)
+    if entry_type == manifest.ENTRY_IS_DIR and not dst_path.endswith('/'):
+      dst_path += '/'
+    entry_info = self.make_zipinfo(path=dst_path, mode=mode)
+
+    if entry_type == manifest.ENTRY_IS_FILE:
+      entry_info.compress_type = zipfile.ZIP_DEFLATED
+      with open(src, 'rb') as src:
+        self.zip_file.writestr(entry_info, src.read())
+    elif entry_type == manifest.ENTRY_IS_DIR:
+      entry_info.compress_type = zipfile.ZIP_STORED
+      # Set directory bits
+      entry_info.external_attr |= (UNIX_DIR_BIT << 16) | MSDOS_DIR_BIT
+      self.zip_file.writestr(entry_info, '')
+    elif entry_type == manifest.ENTRY_IS_LINK:
+      entry_info.compress_type = zipfile.ZIP_STORED
+      # Set directory bits
+      entry_info.external_attr |= (UNIX_SYMLINK_BIT << 16)
+      self.zip_file.writestr(entry_info, src)
+    elif entry_type == manifest.ENTRY_IS_TREE:
+      self.add_tree(src, dst_path, mode)
+    elif entry.entry_type == manifest.ENTRY_IS_EMPTY_FILE:
+      entry_info.compress_type = zipfile.ZIP_DEFLATED
+      self.zip_file.writestr(entry_info, '')
+    else:
+      raise Exception('Unknown type for manifest entry:', entry)
+
+  def add_tree(self, tree_top: str, destpath: str, mode: int):
+    """Add a tree artifact to the zip file.
+
+    Args:
+      tree_top: the top of the tree to add
+      destpath: the path under which to place the files
+      mode: if not None, file mode to apply to all files
+    """
+
+    # We expect /-style paths.
+    tree_top = os.path.normpath(tree_top).replace(os.path.sep, '/')
+
+    # Again, we expect /-style paths.
+    dest = destpath.strip('/')  # redundant, dests should never have / here
+    dest = os.path.normpath(dest).replace(os.path.sep, '/')
+    # paths should not have a leading ./
+    dest = '' if dest == '.' else dest + '/'
+
+    to_write = {}
+    for root, dirs, files in os.walk(tree_top):
+      # While `tree_top` uses '/' as a path separator, results returned by
+      # `os.walk` and `os.path.join` on Windows may not.
+      root = os.path.normpath(root).replace(os.path.sep, '/')
+
+      rel_path_from_top = root[len(tree_top):].lstrip('/')
+      if rel_path_from_top:
+        dest_dir = dest + rel_path_from_top + '/'
+      else:
+        dest_dir = dest
+      for file in files:
+        to_write[dest_dir + file] = root + '/' + file
+
+    for path in sorted(to_write.keys()):
+      content_path = to_write[path]
+      # If mode is unspecified, derive the mode from the file's mode.
+      if mode is None:
+        f_mode = 0o755 if os.access(content_path, os.X_OK) else 0o644
+      else:
+        f_mode = mode
+      entry_info = self.make_zipinfo(path=path, mode=f_mode)
+      entry_info.compress_type = zipfile.ZIP_DEFLATED
+      if not content_path:
+        self.zip_file.writestr(entry_info, '')
+      else:
+        with open(content_path, 'rb') as src:
+          self.zip_file.writestr(entry_info, src.read())
+
 
 def main(args):
   unix_ts = max(ZIP_EPOCH, args.timestamp)
@@ -125,12 +217,14 @@ def main(args):
   if args.mode:
     default_mode = int(args.mode, 8)
 
-  with zipfile.ZipFile(args.output, mode='w') as zip_file:
-    if args.manifest:
-      with open(args.manifest, 'r') as manifest_fp:
-        manifest = json.load(manifest_fp)
-        for entry in manifest:
-          _add_manifest_entry(args, zip_file, entry, default_mode, ts)
+  if not args.manifest:
+    raise Exception('Missing --manifest')
+  with open(args.manifest, 'r') as manifest_fp:
+    manifest = json.load(manifest_fp)
+    with ZipWriter(
+        args.output, time_stamp=ts, default_mode=default_mode) as zip_out:
+      for entry in manifest:
+        zip_out.add_manifest_entry(args, entry)
 
 
 if __name__ == '__main__':

tests/zip/BUILD

@@ -116,6 +116,11 @@ pkg_zip(
     timestamp = 1234567890,
 )
 
+pkg_zip(
+    name = "test_zip_tree",
+    srcs = [":generate_tree"],
+)
+
 pkg_zip(
     name = "test_zip_permissions",
     srcs = [
@@ -206,8 +211,9 @@ py_test(
         # it with this name.
         "test_zip_out.foo",
         "test_zip_package_dir0.zip",
-        "test_zip_timestamp.zip",
         "test_zip_permissions.zip",
+        "test_zip_timestamp.zip",
+        "test_zip_tree",
         "test-zip-strip_prefix-empty.zip",
         "test-zip-strip_prefix-none.zip",
         "test-zip-strip_prefix-zipcontent.zip",

tests/zip/zip_test.py

@@ -77,8 +77,9 @@ def assertZipFileContent(self, zip_file, content):
           expect_dir_bits = UNIX_DIR_BIT << 16 | MSDOS_DIR_BIT
           self.assertEqual(info.external_attr & expect_dir_bits,
                            expect_dir_bits)
-        self.assertEqual((info.external_attr >> 16) & UNIX_RWX_BITS,
-                         expected.get("attr", 0o555))
+        got = (info.external_attr >> 16) & UNIX_RWX_BITS
+        exp = expected.get("attr", 0o555)
+        self.assertEqual(got, exp, 'got %o, expected %o' % (got, exp))
 
   def test_empty(self):
     self.assertZipFileContent("test_zip_empty.zip", [])
@@ -132,6 +133,15 @@ def test_zip_strip_prefix_dot(self):
         {"filename": "zipcontent/loremipsum.txt", "crc": LOREM_CRC},
     ])
 
+  def test_zip_tree(self):
+    self.assertZipFileContent("test_zip_tree.zip", [
+        {"filename": "generate_tree/a/a"},
+        {"filename": "generate_tree/a/b/c"},
+        {"filename": "generate_tree/b/c/d"},
+        {"filename": "generate_tree/b/d"},
+        {"filename": "generate_tree/b/e"},
+    ])
+
 
 class ZipEquivalency(ZipTest):
   """Check that some generated zip files are equivalent to each-other."""