Improve load speed for python payload script

When using update_payload module to read a full OTA, it will attempt to
load the entire payload into memory first, which is slow. Use mmap
instead, as parsing manifest doesn't require the entire payload to be in
memory, only the first couple of MB.
In case of opening a zipfile, use zfp.open(), which returns a streaming
interface.

Test: import update_payload, load an OTA
Change-Id: I3f7c7515dbf5576ccfda8d844fe22e9be34fdaed

1 files changed, 5 insertions, 3 deletions

diff --git a/scripts/update_payload/payload.py b/scripts/update_payload/payload.py
index 6b28046e..998703ad 100644
--- a/scripts/update_payload/payload.py
+++ b/scripts/update_payload/payload.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 
 import hashlib
 import io
+import mmap
 import struct
 import zipfile
 
@@ -123,10 +124,11 @@ class Payload(object):
     """
     if zipfile.is_zipfile(payload_file):
       with zipfile.ZipFile(payload_file) as zfp:
-        with zfp.open("payload.bin") as payload_fp:
-          self.payload_file = io.BytesIO(payload_fp.read())
+        self.payload_file = zfp.open("payload.bin", "r")
     elif isinstance(payload_file, str):
-      self.payload_file = open(payload_file, "rb")
+      payload_fp = open(payload_file, "rb")
+      payload_bytes = mmap.mmap(payload_fp.fileno(), 0, access=mmap.ACCESS_READ)
+      self.payload_file = io.BytesIO(payload_bytes)
     else:
       self.payload_file = payload_file
     self.payload_file_offset = payload_file_offset