sparse_img.py --get_partition_size return size of partition

Also, move code from build_super_image.py to sparse_img.py.

Test: sparse_img.py on sparse and non-sparse images
Bug: 122377935
Change-Id: Ie91fdfdbb54298ea27eb20d1b5363aeb1470356e
Merged-In: Ie91fdfdbb54298ea27eb20d1b5363aeb1470356e

tools/releasetools/build_super_image.py

@@ -58,16 +58,8 @@ logger = logging.getLogger(__name__)
 UNZIP_PATTERN = ["IMAGES/*", "META/*"]
 
 
-def GetPartitionSizeFromImage(img):
-  try:
-    simg = sparse_img.SparseImage(img)
-    return simg.blocksize * simg.total_blocks
-  except ValueError:
-    return os.path.getsize(img)
-
-
 def GetArgumentsForImage(partition, group, image=None):
-  image_size = GetPartitionSizeFromImage(image) if image else 0
+  image_size = sparse_img.GetImagePartitionSize(image) if image else 0
 
   cmd = ["--partition",
          "{}:readonly:{}:{}".format(partition, image_size, group)]

tools/releasetools/sparse_img.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+#
 # Copyright (C) 2014 The Android Open Source Project
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,6 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import print_function
+
+import argparse
 import bisect
 import logging
 import os
@@ -344,3 +349,21 @@ class SparseImage(object):
     """Throw away the file map and treat the entire image as
     undifferentiated data."""
     self.file_map = {"__DATA": self.care_map}
+
+
+def GetImagePartitionSize(img):
+  try:
+    simg = SparseImage(img, build_map=False)
+    return simg.blocksize * simg.total_blocks
+  except ValueError:
+    return os.path.getsize(img)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument('image')
+  parser.add_argument('--get_partition_size', action='store_true',
+                      help='Return partition size of the image')
+  args = parser.parse_args()
+  if args.get_partition_size:
+    print(GetImagePartitionSize(args.image))