diff options
| -rw-r--r-- | gemato/hash.py | 18 | ||||
| -rw-r--r-- | tests/test_hash.py | 12 |
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) |