Merge "releasetools: Don't return ZipFile from common.UnzipTemp()."

tools/releasetools/add_img_to_target_files.py

@@ -632,23 +632,26 @@ def AddImagesToTargetFiles(filename):
   The images will be created under IMAGES/ in the input target_files.zip.
 
   Args:
-      filename: the target_files.zip, or the zip root directory.
+    filename: the target_files.zip, or the zip root directory.
   """
   if os.path.isdir(filename):
     OPTIONS.input_tmp = os.path.abspath(filename)
-    input_zip = None
   else:
-    OPTIONS.input_tmp, input_zip = common.UnzipTemp(filename)
+    OPTIONS.input_tmp = common.UnzipTemp(filename)
 
   if not OPTIONS.add_missing:
     if os.path.isdir(os.path.join(OPTIONS.input_tmp, "IMAGES")):
       print("target_files appears to already contain images.")
       sys.exit(1)
 
-  # {vendor,product}.img is unlike system.img or system_other.img. Because it could
-  # be built from source, or dropped into target_files.zip as a prebuilt blob.
-  # We consider either of them as {vendor,product}.img being available, which could
-  # be used when generating vbmeta.img for AVB.
+  OPTIONS.info_dict = common.LoadInfoDict(OPTIONS.input_tmp, OPTIONS.input_tmp)
+
+  has_recovery = OPTIONS.info_dict.get("no_recovery") != "true"
+
+  # {vendor,product}.img is unlike system.img or system_other.img. Because it
+  # could be built from source, or dropped into target_files.zip as a prebuilt
+  # blob. We consider either of them as {vendor,product}.img being available,
+  # which could be used when generating vbmeta.img for AVB.
   has_vendor = (os.path.isdir(os.path.join(OPTIONS.input_tmp, "VENDOR")) or
                 os.path.exists(os.path.join(OPTIONS.input_tmp, "IMAGES",
                                             "vendor.img")))
@@ -658,16 +661,14 @@ def AddImagesToTargetFiles(filename):
   has_system_other = os.path.isdir(os.path.join(OPTIONS.input_tmp,
                                                 "SYSTEM_OTHER"))
 
-  if input_zip:
-    OPTIONS.info_dict = common.LoadInfoDict(input_zip, OPTIONS.input_tmp)
-
-    common.ZipClose(input_zip)
+  # Set up the output destination. It writes to the given directory for dir
+  # mode; otherwise appends to the given ZIP.
+  if os.path.isdir(filename):
+    output_zip = None
+  else:
     output_zip = zipfile.ZipFile(filename, "a",
                                  compression=zipfile.ZIP_DEFLATED,
                                  allowZip64=True)
-  else:
-    OPTIONS.info_dict = common.LoadInfoDict(filename, filename)
-    output_zip = None
 
   # Always make input_tmp/IMAGES available, since we may stage boot / recovery
   # images there even under zip mode. The directory will be cleaned up as part
@@ -676,8 +677,6 @@ def AddImagesToTargetFiles(filename):
   if not os.path.isdir(images_dir):
     os.makedirs(images_dir)
 
-  has_recovery = (OPTIONS.info_dict.get("no_recovery") != "true")
-
   # A map between partition names and their paths, which could be used when
   # generating AVB vbmeta image.
   partitions = dict()

tools/releasetools/check_target_files_signatures.py

@@ -248,7 +248,7 @@ class TargetFiles(object):
     if compressed_extension:
       apk_extensions.append("*.apk" + compressed_extension)
 
-    d, z = common.UnzipTemp(filename, apk_extensions)
+    d = common.UnzipTemp(filename, apk_extensions)
     try:
       self.apks = {}
       self.apks_by_basename = {}
@@ -283,8 +283,6 @@ class TargetFiles(object):
     finally:
       shutil.rmtree(d)
 
-    z.close()
-
   def CheckSharedUids(self):
     """Look for any instances where packages signed with different
     certs request the same sharedUserId."""

tools/releasetools/common.py

@@ -584,8 +584,7 @@ def UnzipTemp(filename, pattern=None):
   then unzip bar.zip into that_dir/BOOTABLE_IMAGES.
 
   Returns:
-    (tempdir, zipobj): tempdir is the name of the temprary directory; zipobj is
-        a zipfile.ZipFile (of the main file), open for reading.
+    The name of the temporary directory.
   """
 
   def unzip_to_dir(filename, dirname):
@@ -607,7 +606,7 @@ def UnzipTemp(filename, pattern=None):
   else:
     unzip_to_dir(filename, tmp)
 
-  return tmp, zipfile.ZipFile(filename, "r")
+  return tmp
 
 
 def GetSparseImage(which, tmpdir, input_zip, allow_shared_blocks):

tools/releasetools/img_from_target_files.py

@@ -71,8 +71,7 @@ def main(argv):
     common.Usage(__doc__)
     sys.exit(1)
 
-  OPTIONS.input_tmp, input_zip = common.UnzipTemp(
-      args[0], ["IMAGES/*", "OTA/*"])
+  OPTIONS.input_tmp = common.UnzipTemp(args[0], ["IMAGES/*", "OTA/*"])
   output_zip = zipfile.ZipFile(args[1], "w", compression=zipfile.ZIP_DEFLATED)
   CopyInfo(output_zip)
 

tools/releasetools/ota_from_target_files.py

@@ -1248,8 +1248,11 @@ def GetTargetFilesZipForSecondaryImages(input_file, skip_postinstall=False):
   target_file = common.MakeTempFile(prefix="targetfiles-", suffix=".zip")
   target_zip = zipfile.ZipFile(target_file, 'w', allowZip64=True)
 
-  input_tmp, input_zip = common.UnzipTemp(input_file, UNZIP_PATTERN)
-  for info in input_zip.infolist():
+  input_tmp = common.UnzipTemp(input_file, UNZIP_PATTERN)
+  with zipfile.ZipFile(input_file, 'r') as input_zip:
+    infolist = input_zip.infolist()
+
+  for info in infolist:
     unzipped_file = os.path.join(input_tmp, *info.filename.split('/'))
     if info.filename == 'IMAGES/system_other.img':
       common.ZipWrite(target_zip, unzipped_file, arcname='IMAGES/system.img')
@@ -1266,7 +1269,6 @@ def GetTargetFilesZipForSecondaryImages(input_file, skip_postinstall=False):
     elif info.filename.startswith(('META/', 'IMAGES/')):
       common.ZipWrite(target_zip, unzipped_file, arcname=info.filename)
 
-  common.ZipClose(input_zip)
   common.ZipClose(target_zip)
 
   return target_file
@@ -1634,11 +1636,9 @@ def main(argv):
 
   if OPTIONS.extracted_input is not None:
     OPTIONS.input_tmp = OPTIONS.extracted_input
-    input_zip = zipfile.ZipFile(args[0], "r")
   else:
     print("unzipping target target-files...")
-    OPTIONS.input_tmp, input_zip = common.UnzipTemp(
-        args[0], UNZIP_PATTERN)
+    OPTIONS.input_tmp = common.UnzipTemp(args[0], UNZIP_PATTERN)
   OPTIONS.target_tmp = OPTIONS.input_tmp
 
   # If the caller explicitly specified the device-specific extensions path via
@@ -1670,16 +1670,17 @@ def main(argv):
 
   # Generate a full OTA.
   if OPTIONS.incremental_source is None:
-    WriteFullOTAPackage(input_zip, output_zip)
+    with zipfile.ZipFile(args[0], 'r') as input_zip:
+      WriteFullOTAPackage(input_zip, output_zip)
 
   # Generate an incremental OTA.
   else:
     print("unzipping source target-files...")
-    OPTIONS.source_tmp, source_zip = common.UnzipTemp(
-        OPTIONS.incremental_source,
-        UNZIP_PATTERN)
-
-    WriteBlockIncrementalOTAPackage(input_zip, source_zip, output_zip)
+    OPTIONS.source_tmp = common.UnzipTemp(
+        OPTIONS.incremental_source, UNZIP_PATTERN)
+    with zipfile.ZipFile(args[0], 'r') as input_zip, \
+        zipfile.ZipFile(OPTIONS.incremental_source, 'r') as source_zip:
+      WriteBlockIncrementalOTAPackage(input_zip, source_zip, output_zip)
 
     if OPTIONS.log_diff:
       with open(OPTIONS.log_diff, 'w') as out_file:
@@ -1687,7 +1688,6 @@ def main(argv):
         target_files_diff.recursiveDiff(
             '', OPTIONS.source_tmp, OPTIONS.input_tmp, out_file)
 
-  common.ZipClose(input_zip)
   common.ZipClose(output_zip)
 
   # Sign the generated zip package unless no_signing is specified.

tools/releasetools/test_common.py

@@ -523,9 +523,9 @@ class CommonUtilsTest(unittest.TestCase):
       target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 8))
       target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
 
-    tempdir, input_zip = common.UnzipTemp(target_files)
-    sparse_image = common.GetSparseImage('system', tempdir, input_zip, False)
-    input_zip.close()
+    tempdir = common.UnzipTemp(target_files)
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      sparse_image = common.GetSparseImage('system', tempdir, input_zip, False)
 
     self.assertDictEqual(
         {
@@ -552,11 +552,11 @@ class CommonUtilsTest(unittest.TestCase):
       target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 8))
       target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
 
-    tempdir, input_zip = common.UnzipTemp(target_files)
-    self.assertRaises(
-        AssertionError, common.GetSparseImage, 'system', tempdir, input_zip,
-        False)
-    input_zip.close()
+    tempdir = common.UnzipTemp(target_files)
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      self.assertRaises(
+          AssertionError, common.GetSparseImage, 'system', tempdir, input_zip,
+          False)
 
   def test_GetSparseImage_sharedBlocks_notAllowed(self):
     """Tests the case of having overlapping blocks but disallowed."""
@@ -574,11 +574,11 @@ class CommonUtilsTest(unittest.TestCase):
       target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 7))
       target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
 
-    tempdir, input_zip = common.UnzipTemp(target_files)
-    self.assertRaises(
-        AssertionError, common.GetSparseImage, 'system', tempdir, input_zip,
-        False)
-    input_zip.close()
+    tempdir = common.UnzipTemp(target_files)
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      self.assertRaises(
+          AssertionError, common.GetSparseImage, 'system', tempdir, input_zip,
+          False)
 
   def test_GetSparseImage_sharedBlocks_allowed(self):
     """Tests the case for target using BOARD_EXT4_SHARE_DUP_BLOCKS := true."""
@@ -597,9 +597,9 @@ class CommonUtilsTest(unittest.TestCase):
       target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 7))
       target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
 
-    tempdir, input_zip = common.UnzipTemp(target_files)
-    sparse_image = common.GetSparseImage('system', tempdir, input_zip, True)
-    input_zip.close()
+    tempdir = common.UnzipTemp(target_files)
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      sparse_image = common.GetSparseImage('system', tempdir, input_zip, True)
 
     self.assertDictEqual(
         {
@@ -638,9 +638,9 @@ class CommonUtilsTest(unittest.TestCase):
       # '/system/file2' has less blocks listed (2) than actual (3).
       target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
 
-    tempdir, input_zip = common.UnzipTemp(target_files)
-    sparse_image = common.GetSparseImage('system', tempdir, input_zip, False)
-    input_zip.close()
+    tempdir = common.UnzipTemp(target_files)
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      sparse_image = common.GetSparseImage('system', tempdir, input_zip, False)
 
     self.assertFalse(sparse_image.file_map['/system/file1'].extra)
     self.assertTrue(sparse_image.file_map['/system/file2'].extra['incomplete'])

tools/releasetools/validate_target_files.py

@@ -192,9 +192,10 @@ def main(argv):
                       datefmt=date_format)
 
   logging.info("Unzipping the input target_files.zip: %s", args[0])
-  input_tmp, input_zip = common.UnzipTemp(args[0])
+  input_tmp = common.UnzipTemp(args[0])
 
-  ValidateFileConsistency(input_zip, input_tmp)
+  with zipfile.ZipFile(args[0], 'r') as input_zip:
+    ValidateFileConsistency(input_zip, input_tmp)
 
   info_dict = common.LoadInfoDict(input_tmp)
   ValidateInstallRecoveryScript(input_tmp, info_dict)