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

In addition to the unzipping work, common.UnzipTemp() kindly bundles an
open ZipFile object as part of the return value. It doesn't look very
helpful to the callers though. It also looks less obvious that the
caller needs to properly close the handle (missing the close here is
benign though). This CL just removes the ZipFile object out of the
return value, and leaves the work to callers.

Test: `m dist` on both of A/B and non-A/B target.
Test: python -m unittest test_add_img_to_target_files
Test: python -m unittest test_common
Test: python -m unittest test_ota_from_target_files
Test: Check the callers to common.UnzipTemp() in code search.
Change-Id: Id47da3fd42a0e76d6ae8851f05780db319ee48cf

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)