Making MD5 hash method take a buffer instead of bytes.

dhermes · dhermes · commit 86b1f597dea7 · 2015-02-12T10:51:27.000-08:00
This is to account for files too large to fit into memory.
diff --git a/gcloud/storage/_helpers.py b/gcloud/storage/_helpers.py
@@ -192,12 +192,28 @@ def _setter(self, value):
     return property(_getter, _setter)
 
 
-def _base64_md5hash(bytes_to_sign):
+def _write_buffer_to_hash(buffer_object, hash_obj, digest_block_size=8192):
+    """Read blocks from a buffer and update a hash with them.
+
+    :type buffer_object: bytes buffer
+    :param buffer_object: Buffer containing bytes used to update a hash object.
+    """
+    block = buffer_object.read(digest_block_size)
+
+    while len(block) > 0:
+        hash_obj.update(block)
+        # Update the block for the next iteration.
+        block = buffer_object.read(digest_block_size)
+
+
+def _base64_md5hash(buffer_object):
     """Get MD5 hash of bytes (as base64).
 
-    :type bytes_to_sign: bytes
-    :param bytes_to_sign: Bytes used to compute an MD5 hash (as base64).
+    :type buffer_object: bytes buffer
+    :param buffer_object: Buffer containing bytes used to compute an MD5
+                          hash (as base64).
     """
-    hash = MD5.new(data=bytes_to_sign)
-    digest_bytes = hash.digest()
+    hash_obj = MD5.new()
+    _write_buffer_to_hash(buffer_object, hash_obj)
+    digest_bytes = hash_obj.digest()
     return base64.b64encode(digest_bytes)
diff --git a/gcloud/storage/test__helpers.py b/gcloud/storage/test__helpers.py
@@ -244,26 +244,44 @@ def _callFUT(self, bytes_to_sign):
         return _base64_md5hash(bytes_to_sign)
 
     def test_it(self):
+        from io import BytesIO
         BYTES_TO_SIGN = b'FOO'
-        SIGNED_CONTENT = self._callFUT(BYTES_TO_SIGN)
+        BUFFER = BytesIO()
+        BUFFER.write(BYTES_TO_SIGN)
+        BUFFER.seek(0)
+
+        SIGNED_CONTENT = self._callFUT(BUFFER)
         self.assertEqual(SIGNED_CONTENT, b'kBiQqOnIz21aGlQrIp/r/w==')
 
     def test_it_with_stubs(self):
         from gcloud._testing import _Monkey
         from gcloud.storage import _helpers as MUT
 
+        class _Buffer(object):
+
+            def __init__(self, return_vals):
+                self.return_vals = return_vals
+                self._block_sizes = []
+
+            def read(self, block_size):
+                self._block_sizes.append(block_size)
+                return self.return_vals.pop()
+
         BASE64 = _Base64()
         DIGEST_VAL = object()
-        BYTES_TO_SIGN = object()
+        BYTES_TO_SIGN = b'BYTES_TO_SIGN'
+        BUFFER = _Buffer([b'', BYTES_TO_SIGN])
         MD5 = _MD5(DIGEST_VAL)
 
         with _Monkey(MUT, base64=BASE64, MD5=MD5):
-            SIGNED_CONTENT = self._callFUT(BYTES_TO_SIGN)
+            SIGNED_CONTENT = self._callFUT(BUFFER)
 
+        self.assertEqual(BUFFER._block_sizes, [8192, 8192])
         self.assertTrue(SIGNED_CONTENT is DIGEST_VAL)
         self.assertEqual(BASE64._called_b64encode, [DIGEST_VAL])
-        self.assertEqual(MD5._new_called, [BYTES_TO_SIGN])
+        self.assertEqual(MD5._new_called, [None])
         self.assertEqual(MD5.hash_obj.num_digest_calls, 1)
+        self.assertEqual(MD5.hash_obj._blocks, [BYTES_TO_SIGN])
 
 
 class _Connection(object):
@@ -283,11 +301,16 @@ class _MD5Hash(object):
     def __init__(self, digest_val):
         self.digest_val = digest_val
         self.num_digest_calls = 0
+        self._blocks = []
+
+    def update(self, block):
+        self._blocks.append(block)
 
     def digest(self):
         self.num_digest_calls += 1
         return self.digest_val
 
+
 class _MD5(object):
 
     def __init__(self, digest_val):
diff --git a/regression/storage.py b/regression/storage.py
@@ -100,7 +100,7 @@ def setUpClass(cls):
         super(TestStorageFiles, cls).setUpClass()
         for file_data in cls.FILES.values():
             with open(file_data['path'], 'rb') as file_obj:
-                file_data['hash'] = _base64_md5hash(file_obj.read())
+                file_data['hash'] = _base64_md5hash(file_obj)
         cls.bucket = SHARED_BUCKETS['test_bucket']
 
     def setUp(self):
