1 files changed, 46 insertions, 7 deletions

diff --git a/golden/extract_lsdump.py b/golden/extract_lsdump.py
index a6d95c7..fd42be6 100755
--- a/golden/extract_lsdump.py
+++ b/golden/extract_lsdump.py
@@ -19,6 +19,7 @@ import argparse
 import gzip
 import json
 import os
+import zipfile
 
 
 class LsdumpError(Exception):
@@ -205,35 +206,62 @@ def _ParseVtablesFromLsdump(lsdump, output_dump):
         output_dump.record_types.append(record_type)
 
 
-def ParseLsdumpFile(input_path, output_path):
-    """Converts an lsdump file to a dump file for the ABI test.
+def _ParseLsdumpFileObject(lsdump_file):
+    """Converts an lsdump file to a dump object.
 
     Args:
-        input_path: The path to the (gzipped) lsdump file.
-        output_path: The path to the output dump file.
+        lsdump_file: The file object of the input lsdump.
+
+    Returns:
+        The AttrDict object converted from the lsdump file.
 
     Raises:
         LsdumpError if fails to create the dump file.
     """
     try:
-        with _OpenFileOrGzipped(input_path) as lsdump_file:
-            lsdump = json.load(lsdump_file, object_hook=AttrDict)
-    except (IOError, ValueError) as e:
+        lsdump = json.load(lsdump_file, object_hook=AttrDict)
+    except ValueError:
         raise LsdumpError(e)
 
     try:
         output_dump = AttrDict()
         _ParseVtablesFromLsdump(lsdump, output_dump)
         _ParseSymbolsFromLsdump(lsdump, output_dump)
+        return output_dump
     except AttributeError as e:
         raise LsdumpError(e)
 
+
+def _ParseLsdumpZipFile(lsdump_zip, output_zip):
+    """Converts zipped lsdump files to the dump files for ABI test."""
+    for name in lsdump_zip.namelist():
+        if name.endswith(".lsdump"):
+            with lsdump_zip.open(name, mode="r") as lsdump_file:
+                output_dump = _ParseLsdumpFileObject(lsdump_file)
+            output_str = json.dumps(output_dump, indent=1,
+                                    separators=(',', ':'))
+            output_zip.writestr(name, output_str)
+
+
+def ParseLsdumpFile(input_path, output_path):
+    """Converts an lsdump file to a dump file for the ABI test.
+
+    Args:
+        input_path: The path to the (gzipped) lsdump file.
+        output_path: The path to the output dump file.
+
+    Raises:
+        LsdumpError if fails to create the dump file.
+    """
     abs_output_path = os.path.abspath(output_path)
     abs_output_dir = os.path.dirname(abs_output_path)
 
     try:
         if abs_output_dir and not os.path.exists(abs_output_dir):
             os.makedirs(abs_output_dir)
+
+        with _OpenFileOrGzipped(input_path) as lsdump_file:
+            output_dump = _ParseLsdumpFileObject(lsdump_file)
         with open(output_path, 'wb') as output_file:
             json.dump(output_dump, output_file,
                       indent=1, separators=(',', ':'))
@@ -251,6 +279,17 @@ def main():
                             help='output dump file path.')
     args = arg_parser.parse_args()
 
+    # TODO(b/150663999): Remove this when the build system is able to process
+    #                    the large file group.
+    if zipfile.is_zipfile(args.input_path):
+        with zipfile.ZipFile(args.input_path, mode='r') as input_zip:
+            # The zip will be added to a Python package. It is not necessary
+            # to reduce the file size.
+            with zipfile.ZipFile(args.output_path, mode='w',
+                                 compression=zipfile.ZIP_STORED) as output_zip:
+                _ParseLsdumpZipFile(input_zip, output_zip)
+        exit(0)
+
     try:
         ParseLsdumpFile(args.input_path, args.output_path)
     except LsdumpError as e: