hash: Add a size-counting hash

2 files changed, 30 insertions, 0 deletions

diff --git a/gemato/hash.py b/gemato/hash.py
index 18313f7..c7bd971 100644
--- a/gemato/hash.py
+++ b/gemato/hash.py
@@ -16,6 +16,21 @@ class UnsupportedHash(Exception):
 				'Unsupported hash name: {}'.format(hash_name))
 
 
+class SizeHash(object):
+	"""
+	A cheap wrapper to count file size via hashlib-like interface.
+	"""
+
+	def __init__(self):
+		self.size = 0
+
+	def update(self, data):
+		self.size += len(data)
+
+	def hexdigest(self):
+		return self.size
+
+
 def get_hash_by_name(name):
 	"""
 	Get a hashlib-compatible hash object for hash named @name. Supports
@@ -26,6 +41,9 @@ def get_hash_by_name(name):
 	except ValueError:
 		pass
 
+	if name == '__size__':
+		return SizeHash()
+
 	# fallback support
 	if name.startswith('sha3_'):
 		try:
diff --git a/tests/test_hash.py b/tests/test_hash.py
index 8ed4617..4ac21ca 100644
--- a/tests/test_hash.py
+++ b/tests/test_hash.py
@@ -245,3 +245,15 @@ class OptionalHashTest(unittest.TestCase):
                     '19fa61d75522a4669b44e39c1d2e1726c530232130d407f89afee0964997f7a73e83be698b288febcf88e3e03c4f0757ea8964e59b63d93708b138cc42a66eb3')
         except gemato.hash.UnsupportedHash:
             raise unittest.SkipTest('hash not supported')
+
+
+class SizeHashTest(unittest.TestCase):
+    """
+    Test __size__ special function.
+    """
+
+    def test_size(self):
+        self.assertEqual(gemato.hash.hash_bytes(TEST_STRING, '__size__'), 43)
+
+    def test_size_empty(self):
+        self.assertEqual(gemato.hash.hash_bytes(b'', '__size__'), 0)