Merge "releasetools: Make common.ZipWriteStr Python 3 compatible."

tools/releasetools/common.py

@@ -863,7 +863,7 @@ def GetUserImage(which, tmpdir, input_zip,
     A Image object. If it is a sparse image and reset_file_map is False, the
     image will have file_map info loaded.
   """
-  if info_dict == None:
+  if info_dict is None:
     info_dict = LoadInfoDict(input_zip)
 
   is_sparse = info_dict.get("extfs_sparse_flag")
@@ -1568,6 +1568,15 @@ def ZipWriteStr(zip_file, zinfo_or_arcname, data, perms=None,
       perms = 0o100644
   else:
     zinfo = zinfo_or_arcname
+    # Python 2 and 3 behave differently when calling ZipFile.writestr() with
+    # zinfo.external_attr being 0. Python 3 uses `0o600 << 16` as the value for
+    # such a case (since
+    # https://github.com/python/cpython/commit/18ee29d0b870caddc0806916ca2c823254f1a1f9),
+    # which seems to make more sense. Otherwise the entry will have 0o000 as the
+    # permission bits. We follow the logic in Python 3 to get consistent
+    # behavior between using the two versions.
+    if not zinfo.external_attr:
+      zinfo.external_attr = 0o600 << 16
 
   # If compress_type is given, it overrides the value in zinfo.
   if compress_type is not None:
@@ -1600,7 +1609,7 @@ def ZipDelete(zip_filename, entries):
   Raises:
     AssertionError: In case of non-zero return from 'zip'.
   """
-  if isinstance(entries, basestring):
+  if isinstance(entries, str):
     entries = [entries]
   cmd = ["zip", "-d", zip_filename] + entries
   RunAndCheckOutput(cmd)

tools/releasetools/test_common.py

@@ -41,7 +41,7 @@ def get_2gb_string():
   # Generate a long string with holes, e.g. 'xyz\x00abc\x00...'.
   for _ in range(0, size, step_size):
     yield os.urandom(block_size)
-    yield '\0' * (step_size - block_size)
+    yield b'\0' * (step_size - block_size)
 
 
 class CommonZipTest(test_utils.ReleaseToolsTestCase):
@@ -72,7 +72,7 @@ class CommonZipTest(test_utils.ReleaseToolsTestCase):
     # Verify the zip contents.
     entry = zip_file.open(arcname)
     sha1_hash = sha1()
-    for chunk in iter(lambda: entry.read(4 * MiB), ''):
+    for chunk in iter(lambda: entry.read(4 * MiB), b''):
       sha1_hash.update(chunk)
     self.assertEqual(expected_hash, sha1_hash.hexdigest())
     self.assertIsNone(zip_file.testzip())
@@ -97,8 +97,8 @@ class CommonZipTest(test_utils.ReleaseToolsTestCase):
     try:
       sha1_hash = sha1()
       for data in contents:
-        sha1_hash.update(data)
-        test_file.write(data)
+        sha1_hash.update(bytes(data))
+        test_file.write(bytes(data))
       test_file.close()
 
       expected_stat = os.stat(test_file_name)
@@ -136,8 +136,11 @@ class CommonZipTest(test_utils.ReleaseToolsTestCase):
         expected_mode = extra_args.get("perms", 0o644)
       else:
         arcname = zinfo_or_arcname.filename
-        expected_mode = extra_args.get("perms",
-                                       zinfo_or_arcname.external_attr >> 16)
+        if zinfo_or_arcname.external_attr:
+          zinfo_perms = zinfo_or_arcname.external_attr >> 16
+        else:
+          zinfo_perms = 0o600
+        expected_mode = extra_args.get("perms", zinfo_perms)
 
       common.ZipWriteStr(zip_file, zinfo_or_arcname, contents, **extra_args)
       common.ZipClose(zip_file)
@@ -262,6 +265,10 @@ class CommonZipTest(test_utils.ReleaseToolsTestCase):
         "perms": 0o600,
         "compress_type": zipfile.ZIP_STORED,
     })
+    self._test_ZipWriteStr(zinfo, random_string, {
+        "perms": 0o000,
+        "compress_type": zipfile.ZIP_STORED,
+    })
 
   def test_ZipWriteStr_large_file(self):
     # zipfile.writestr() doesn't work when the str size is over 2GiB even with
@@ -274,9 +281,9 @@ class CommonZipTest(test_utils.ReleaseToolsTestCase):
     })
 
   def test_ZipWriteStr_resets_ZIP64_LIMIT(self):
-    self._test_reset_ZIP64_LIMIT(self._test_ZipWriteStr, "foo", "")
+    self._test_reset_ZIP64_LIMIT(self._test_ZipWriteStr, 'foo', b'')
     zinfo = zipfile.ZipInfo(filename="foo")
-    self._test_reset_ZIP64_LIMIT(self._test_ZipWriteStr, zinfo, "")
+    self._test_reset_ZIP64_LIMIT(self._test_ZipWriteStr, zinfo, b'')
 
   def test_bug21309935(self):
     zip_file = tempfile.NamedTemporaryFile(delete=False)