Add assets support to aar_import.

Fixes https://github.com/bazelbuild/bazel/issues/4439.

Empty assets are handled by writing out an empty file named "assets/empty_asset_generated_by_bazel~", which will silently be ignored by AAPT.

RELNOTES: aar_import now supports assets.
PiperOrigin-RevId: 182110400

2 files changed, 92 insertions, 36 deletions

diff --git a/tools/android/aar_resources_extractor.py b/tools/android/aar_resources_extractor.py
index 5d875fbe12..1a18e24402 100644
--- a/tools/android/aar_resources_extractor.py
+++ b/tools/android/aar_resources_extractor.py
@@ -35,6 +35,7 @@ gflags.DEFINE_string("input_aar", None, "Input AAR")
 gflags.MarkFlagAsRequired("input_aar")
 gflags.DEFINE_string("output_res_dir", None, "Output resources directory")
 gflags.MarkFlagAsRequired("output_res_dir")
+gflags.DEFINE_string("output_assets_dir", None, "Output assets directory")
 
 
 def ExtractResources(aar, output_res_dir):
@@ -43,51 +44,77 @@ def ExtractResources(aar, output_res_dir):
   output_res_dir_abs = os.path.abspath(output_res_dir)
   for name in aar.namelist():
     if name.startswith("res/"):
-      if os.name == "nt":
-        fullpath = os.path.normpath(os.path.join(output_res_dir_abs, name))
-        if name[-1] == "/":
-          # The zip entry is a directory. Create a junction to it, which also
-          # takes care of creating the directory and all of its parents in a
-          # longpath-safe manner.
-          # We must pretend to have extracted this directory, even if it's
-          # empty, therefore we mustn't rely on creating it as a parent
-          # directory of a subsequently extracted zip entry (because there may
-          # be no such subsequent entry).
-          with junction.TempJunction(fullpath.rstrip("/")) as juncpath:
-            pass
-        else:
-          # The zip entry is a file. Create a junction to its parent directory,
-          # then open the compressed entry as a file object, so we can extract
-          # the data even if the extracted file's path would be too long.
-          # The tradeoff is that we lose the permission bits of the compressed
-          # file, but Unix permissions don't mean much on Windows anyway.
-          with junction.TempJunction(os.path.dirname(fullpath)) as juncpath:
-            extracted_path = os.path.join(juncpath, os.path.basename(fullpath))
-            with aar.open(name) as src_fd:
-              with open(extracted_path, "wb") as dest_fd:
-                dest_fd.write(src_fd.read())
-      else:
-        aar.extract(name, output_res_dir)
+      ExtractOneFile(aar, name, output_res_dir_abs)
       aar_contains_no_resources = False
-
   if aar_contains_no_resources:
     empty_xml_filename = output_res_dir + "/res/values/empty.xml"
-    if os.name == "nt":
-      # Create a junction to the parent directory, because its path might be too
-      # long. Creating the junction also creates all parent directories.
-      with junction.TempJunction(os.path.dirname(empty_xml_filename)) as junc:
-        xmlpath = os.path.join(junc, os.path.basename(empty_xml_filename))
-        with open(xmlpath, "wb") as empty_xml:
-          empty_xml.write(b"<resources/>")
+    WriteFileWithJunctions(empty_xml_filename, b"<resources/>")
+
+
+def ExtractAssets(aar, output_assets_dir):
+  """Extracts assets from an `aar` file to the `output_assets_dir` directory."""
+  aar_contains_no_assets = True
+  output_assets_dir_abs = os.path.abspath(output_assets_dir)
+  for name in aar.namelist():
+    if name.startswith("assets/"):
+      ExtractOneFile(aar, name, output_assets_dir_abs)
+      aar_contains_no_assets = False
+  if aar_contains_no_assets:
+    # aapt will ignore this file and not print an error message, because it
+    # thinks that it is a swap file. We need to create at least one file so that
+    # Bazel does not complain that the output tree artifact was not created.
+    empty_asset_filename = (output_assets_dir +
+                            "/assets/empty_asset_generated_by_bazel~")
+    WriteFileWithJunctions(empty_asset_filename, b"")
+
+
+def WriteFileWithJunctions(filename, content):
+  """Writes file including creating any junctions or directories necessary."""
+  if os.name == "nt":
+    # Create a junction to the parent directory, because its path might be too
+    # long. Creating the junction also creates all parent directories.
+    with junction.TempJunction(os.path.dirname(filename)) as junc:
+      filename = os.path.join(junc, os.path.basename(filename))
+  else:
+    os.makedirs(os.path.dirname(filename))
+  with open(filename, "wb") as openfile:
+    openfile.write(content)
+
+
+def ExtractOneFile(aar, name, abs_output_dir):
+  """Extract one file from the aar to the output directory."""
+  if os.name == "nt":
+    fullpath = os.path.normpath(os.path.join(abs_output_dir, name))
+    if name[-1] == "/":
+      # The zip entry is a directory. Create a junction to it, which also
+      # takes care of creating the directory and all of its parents in a
+      # longpath-safe manner.
+      # We must pretend to have extracted this directory, even if it's
+      # empty, therefore we mustn't rely on creating it as a parent
+      # directory of a subsequently extracted zip entry (because there may
+      # be no such subsequent entry).
+      with junction.TempJunction(fullpath.rstrip("/")) as juncpath:
+        pass
     else:
-      os.makedirs(os.path.dirname(empty_xml_filename))
-      with open(empty_xml_filename, "wb") as empty_xml:
-        empty_xml.write(b"<resources/>")
+      # The zip entry is a file. Create a junction to its parent directory,
+      # then open the compressed entry as a file object, so we can extract
+      # the data even if the extracted file's path would be too long.
+      # The tradeoff is that we lose the permission bits of the compressed
+      # file, but Unix permissions don't mean much on Windows anyway.
+      with junction.TempJunction(os.path.dirname(fullpath)) as juncpath:
+        extracted_path = os.path.join(juncpath, os.path.basename(fullpath))
+        with aar.open(name) as src_fd:
+          with open(extracted_path, "wb") as dest_fd:
+            dest_fd.write(src_fd.read())
+  else:
+    aar.extract(name, abs_output_dir)
 
 
 def main():
   with zipfile.ZipFile(FLAGS.input_aar, "r") as aar:
     ExtractResources(aar, FLAGS.output_res_dir)
+    if FLAGS.output_assets_dir is not None:
+      ExtractAssets(aar, FLAGS.output_assets_dir)
 
 if __name__ == "__main__":
   FLAGS(sys.argv)
diff --git a/tools/android/aar_resources_extractor_test.py b/tools/android/aar_resources_extractor_test.py
index ba93c320b9..18bb5b1daa 100644
--- a/tools/android/aar_resources_extractor_test.py
+++ b/tools/android/aar_resources_extractor_test.py
@@ -61,6 +61,7 @@ class AarResourcesExtractorTest(unittest.TestCase):
     aar = zipfile.ZipFile(io.BytesIO(), "w")
     aar.writestr("res/values/values.xml", "some values")
     aar.writestr("res/layouts/layout.xml", "some layout")
+    aar.writestr("assets/a", "some asset")
     os.makedirs("out_dir")
     aar_resources_extractor.ExtractResources(aar, "out_dir")
     expected_resources = [
@@ -73,6 +74,34 @@ class AarResourcesExtractorTest(unittest.TestCase):
     with open("out_dir/res/layouts/layout.xml", "r") as layout_xml:
       self.assertEqual("some layout", layout_xml.read())
 
+  def testNoAssets(self):
+    aar = zipfile.ZipFile(io.BytesIO(), "w")
+    os.makedirs("out_dir")
+    aar_resources_extractor.ExtractAssets(aar, "out_dir")
+    expected_assets = [
+        _HostPath("out_dir/assets/empty_asset_generated_by_bazel~")
+    ]
+    self.assertEqual(expected_assets, self.DirContents("out_dir"))
+    self.assertEqual(
+        os.stat("out_dir/assets/empty_asset_generated_by_bazel~").st_size, 0)
+
+  def testContainsAssets(self):
+    aar = zipfile.ZipFile(io.BytesIO(), "w")
+    aar.writestr("res/values/values.xml", "some values")
+    aar.writestr("assets/a", "some asset")
+    aar.writestr("assets/b", "some other asset")
+    os.makedirs("out_dir")
+    aar_resources_extractor.ExtractAssets(aar, "out_dir")
+    expected_resources = [
+        _HostPath("out_dir/assets/a"),
+        _HostPath("out_dir/assets/b")
+    ]
+    self.assertCountEqual(expected_resources, self.DirContents("out_dir"))
+    with open("out_dir/assets/a", "r") as values_xml:
+      self.assertEqual("some asset", values_xml.read())
+    with open("out_dir/assets/b", "r") as layout_xml:
+      self.assertEqual("some other asset", layout_xml.read())
+
 
 if __name__ == "__main__":
   unittest.main()