Support merging target files from directory

Expand `merge_target_files.py` API capabilities so that
`--framework-target-files` and `--vendor-target-files`
can be either zip archives or directories.

Test: Create a merged package by vendor target files folder
Test: atest --host releasetools_test
Bug: 276068400
Change-Id: I200be2a458ae59a61e05bfd7c78ab66093db32eb

tools/releasetools/merge/merge_meta.py

@@ -99,16 +99,16 @@ def MergeMetaFiles(temp_dir, merged_dir):
   """Merges various files in META/*."""
 
   framework_meta_dir = os.path.join(temp_dir, 'framework_meta', 'META')
-  merge_utils.ExtractItems(
+  merge_utils.CollectTargetFiles(
-      input_zip=OPTIONS.framework_target_files,
+      input_zipfile_or_dir=OPTIONS.framework_target_files,
       output_dir=os.path.dirname(framework_meta_dir),
-      extract_item_list=('META/*',))
+      item_list=('META/*',))
 
   vendor_meta_dir = os.path.join(temp_dir, 'vendor_meta', 'META')
-  merge_utils.ExtractItems(
+  merge_utils.CollectTargetFiles(
-      input_zip=OPTIONS.vendor_target_files,
+      input_zipfile_or_dir=OPTIONS.vendor_target_files,
       output_dir=os.path.dirname(vendor_meta_dir),
-      extract_item_list=('META/*',))
+      item_list=('META/*',))
 
   merged_meta_dir = os.path.join(merged_dir, 'META')
 

tools/releasetools/merge/merge_target_files.py

@@ -26,9 +26,9 @@ This script produces a complete, merged target files package:
 
 Usage: merge_target_files [args]
 
-  --framework-target-files framework-target-files-zip-archive
+  --framework-target-files framework-target-files-package
       The input target files package containing framework bits. This is a zip
-      archive.
+      archive or a directory.
 
   --framework-item-list framework-item-list-file
       The optional path to a newline-separated config file of items that
@@ -38,9 +38,9 @@ Usage: merge_target_files [args]
       The optional path to a newline-separated config file of keys to
       extract from the framework META/misc_info.txt file.
 
-  --vendor-target-files vendor-target-files-zip-archive
+  --vendor-target-files vendor-target-files-package
       The input target files package containing vendor bits. This is a zip
-      archive.
+      archive or a directory.
 
   --vendor-item-list vendor-item-list-file
       The optional path to a newline-separated config file of items that
@@ -172,18 +172,18 @@ def create_merged_package(temp_dir):
     Path to merged package under temp directory.
   """
   # Extract "as is" items from the input framework and vendor partial target
-  # files packages directly into the output temporary directory, since these items
+  # files packages directly into the output temporary directory, since these
-  # do not need special case processing.
+  # items do not need special case processing.
 
   output_target_files_temp_dir = os.path.join(temp_dir, 'output')
-  merge_utils.ExtractItems(
+  merge_utils.CollectTargetFiles(
-      input_zip=OPTIONS.framework_target_files,
+      input_zipfile_or_dir=OPTIONS.framework_target_files,
       output_dir=output_target_files_temp_dir,
-      extract_item_list=OPTIONS.framework_item_list)
+      item_list=OPTIONS.framework_item_list)
-  merge_utils.ExtractItems(
+  merge_utils.CollectTargetFiles(
-      input_zip=OPTIONS.vendor_target_files,
+      input_zipfile_or_dir=OPTIONS.vendor_target_files,
       output_dir=output_target_files_temp_dir,
-      extract_item_list=OPTIONS.vendor_item_list)
+      item_list=OPTIONS.vendor_item_list)
 
   # Perform special case processing on META/* items.
   # After this function completes successfully, all the files we need to create
@@ -231,7 +231,8 @@ def rebuild_image_with_sepolicy(target_files_dir):
   def copy_selinux_file(input_path, output_filename):
     input_filename = os.path.join(target_files_dir, input_path)
     if not os.path.exists(input_filename):
-      input_filename = input_filename.replace('SYSTEM_EXT/', 'SYSTEM/system_ext/') \
+      input_filename = input_filename.replace('SYSTEM_EXT/',
+                                              'SYSTEM/system_ext/') \
           .replace('PRODUCT/', 'SYSTEM/product/')
       if not os.path.exists(input_filename):
         logger.info('Skipping copy_selinux_file for %s', input_filename)
@@ -272,7 +273,10 @@ def rebuild_image_with_sepolicy(target_files_dir):
   vendor_target_files_dir = common.MakeTempDir(
       prefix='merge_target_files_vendor_target_files_')
   common.UnzipToDir(OPTIONS.vendor_otatools, vendor_otatools_dir)
-  common.UnzipToDir(OPTIONS.vendor_target_files, vendor_target_files_dir)
+  merge_utils.CollectTargetFiles(
+      input_zipfile_or_dir=OPTIONS.vendor_target_files,
+      output_dir=vendor_target_files_dir,
+      item_list=OPTIONS.vendor_item_list)
 
   # Copy the partition contents from the merged target-files archive to the
   # vendor target-files archive.
@@ -303,8 +307,9 @@ def rebuild_image_with_sepolicy(target_files_dir):
   shutil.move(
       os.path.join(vendor_target_files_dir, 'IMAGES', partition_img),
       os.path.join(target_files_dir, 'IMAGES', partition_img))
-  move_only_exists(os.path.join(vendor_target_files_dir, 'IMAGES', partition_map),
+  move_only_exists(
-        os.path.join(target_files_dir, 'IMAGES', partition_map))
+      os.path.join(vendor_target_files_dir, 'IMAGES', partition_map),
+      os.path.join(target_files_dir, 'IMAGES', partition_map))
 
   def copy_recovery_file(filename):
     for subdir in ('VENDOR', 'SYSTEM/vendor'):
@@ -578,10 +583,10 @@ def main():
     common.Usage(__doc__)
     sys.exit(1)
 
-  with zipfile.ZipFile(OPTIONS.framework_target_files, allowZip64=True) as fz:
+  framework_namelist = merge_utils.GetTargetFilesItems(
-    framework_namelist = fz.namelist()
+      OPTIONS.framework_target_files)
-  with zipfile.ZipFile(OPTIONS.vendor_target_files, allowZip64=True) as vz:
+  vendor_namelist = merge_utils.GetTargetFilesItems(
-    vendor_namelist = vz.namelist()
+      OPTIONS.vendor_target_files)
 
   if OPTIONS.framework_item_list:
     OPTIONS.framework_item_list = common.LoadListFromFile(

tools/releasetools/merge/merge_utils.py

@@ -49,28 +49,80 @@ def ExtractItems(input_zip, output_dir, extract_item_list):
   common.UnzipToDir(input_zip, output_dir, filtered_extract_item_list)
 
 
-def CopyItems(from_dir, to_dir, patterns):
+def CopyItems(from_dir, to_dir, copy_item_list):
-  """Similar to ExtractItems() except uses an input dir instead of zip."""
+  """Copies the items in copy_item_list from source to destination directory.
-  file_paths = []
-  for dirpath, _, filenames in os.walk(from_dir):
-    file_paths.extend(
-        os.path.relpath(path=os.path.join(dirpath, filename), start=from_dir)
-        for filename in filenames)
 
-  filtered_file_paths = set()
+  copy_item_list may include files and directories. Will copy the matched
-  for pattern in patterns:
+  files and create the matched directories.
-    filtered_file_paths.update(fnmatch.filter(file_paths, pattern))
 
-  for file_path in filtered_file_paths:
+  Args:
-    original_file_path = os.path.join(from_dir, file_path)
+    from_dir: The source directory.
-    copied_file_path = os.path.join(to_dir, file_path)
+    to_dir: The destination directory.
-    copied_file_dir = os.path.dirname(copied_file_path)
+    copy_item_list: Items to be copied.
-    if not os.path.exists(copied_file_dir):
+  """
-      os.makedirs(copied_file_dir)
+  item_paths = []
-    if os.path.islink(original_file_path):
+  for root, dirs, files in os.walk(from_dir):
-      os.symlink(os.readlink(original_file_path), copied_file_path)
+    item_paths.extend(
+        os.path.relpath(path=os.path.join(root, item_name), start=from_dir)
+        for item_name in files + dirs)
+
+  filtered = set()
+  for pattern in copy_item_list:
+    filtered.update(fnmatch.filter(item_paths, pattern))
+
+  for item in filtered:
+    original_path = os.path.join(from_dir, item)
+    copied_path = os.path.join(to_dir, item)
+    copied_parent_path = os.path.dirname(copied_path)
+    if not os.path.exists(copied_parent_path):
+      os.makedirs(copied_parent_path)
+    if os.path.islink(original_path):
+      os.symlink(os.readlink(original_path), copied_path)
+    elif os.path.isdir(original_path):
+      if not os.path.exists(copied_path):
+        os.makedirs(copied_path)
     else:
-      shutil.copyfile(original_file_path, copied_file_path)
+      shutil.copyfile(original_path, copied_path)
+
+
+def GetTargetFilesItems(target_files_zipfile_or_dir):
+  """Gets a list of target files items."""
+  if zipfile.is_zipfile(target_files_zipfile_or_dir):
+    with zipfile.ZipFile(target_files_zipfile_or_dir, allowZip64=True) as fz:
+      return fz.namelist()
+  elif os.path.isdir(target_files_zipfile_or_dir):
+    item_list = []
+    for root, dirs, files in os.walk(target_files_zipfile_or_dir):
+      item_list.extend(
+          os.path.relpath(path=os.path.join(root, item),
+                          start=target_files_zipfile_or_dir)
+          for item in dirs + files)
+    return item_list
+  else:
+    raise ValueError('Target files should be either zipfile or directory.')
+
+
+def CollectTargetFiles(input_zipfile_or_dir, output_dir, item_list=None):
+  """Extracts input zipfile or copy input directory to output directory.
+
+  Extracts the input zipfile if `input_zipfile_or_dir` is a zip archive, or
+  copies the items if `input_zipfile_or_dir` is a directory.
+
+  Args:
+    input_zipfile_or_dir: The input target files, could be either a zipfile to
+      extract or a directory to copy.
+    output_dir: The output directory that the input files are either extracted
+      or copied.
+    item_list: Files to be extracted or copied. Will extract or copy all files
+      if omitted.
+  """
+  patterns = item_list if item_list else ('*',)
+  if zipfile.is_zipfile(input_zipfile_or_dir):
+    ExtractItems(input_zipfile_or_dir, output_dir, patterns)
+  elif os.path.isdir(input_zipfile_or_dir):
+    CopyItems(input_zipfile_or_dir, output_dir, patterns)
+  else:
+    raise ValueError('Target files should be either zipfile or directory.')
 
 
 def WriteSortedData(data, path):

tools/releasetools/merge/test_merge_utils.py

@@ -35,22 +35,27 @@ class MergeUtilsTest(test_utils.ReleaseToolsTestCase):
       open(path, 'a').close()
       return path
 
+    def createEmptyFolder(path):
+      os.makedirs(path)
+      return path
+
     def createSymLink(source, dest):
       os.symlink(source, dest)
       return dest
 
     def getRelPaths(start, filepaths):
       return set(
-          os.path.relpath(path=filepath, start=start) for filepath in filepaths)
+          os.path.relpath(path=filepath, start=start)
+          for filepath in filepaths)
 
     input_dir = common.MakeTempDir()
     output_dir = common.MakeTempDir()
     expected_copied_items = []
     actual_copied_items = []
-    patterns = ['*.cpp', 'subdir/*.txt']
+    patterns = ['*.cpp', 'subdir/*.txt', 'subdir/empty_dir']
 
-    # Create various files that we expect to get copied because they
+    # Create various files and empty directories that we expect to get copied
-    # match one of the patterns.
+    # because they match one of the patterns.
     expected_copied_items.extend([
         createEmptyFile(os.path.join(input_dir, 'a.cpp')),
         createEmptyFile(os.path.join(input_dir, 'b.cpp')),
@@ -58,6 +63,7 @@ class MergeUtilsTest(test_utils.ReleaseToolsTestCase):
         createEmptyFile(os.path.join(input_dir, 'subdir', 'd.txt')),
         createEmptyFile(
             os.path.join(input_dir, 'subdir', 'subsubdir', 'e.txt')),
+        createEmptyFolder(os.path.join(input_dir, 'subdir', 'empty_dir')),
         createSymLink('a.cpp', os.path.join(input_dir, 'a_link.cpp')),
     ])
     # Create some more files that we expect to not get copied.
@@ -70,9 +76,13 @@ class MergeUtilsTest(test_utils.ReleaseToolsTestCase):
     merge_utils.CopyItems(input_dir, output_dir, patterns)
 
     # Assert the actual copied items match the ones we expected.
-    for dirpath, _, filenames in os.walk(output_dir):
+    for root_dir, dirs, files in os.walk(output_dir):
       actual_copied_items.extend(
-          os.path.join(dirpath, filename) for filename in filenames)
+          os.path.join(root_dir, filename) for filename in files)
+      for dirname in dirs:
+        dir_path = os.path.join(root_dir, dirname)
+        if not os.listdir(dir_path):
+          actual_copied_items.append(dir_path)
     self.assertEqual(
         getRelPaths(output_dir, actual_copied_items),
         getRelPaths(input_dir, expected_copied_items))