cleaned up ChecksummedType

diff --git a/p2pool/bitcoin/data.py b/p2pool/bitcoin/data.py
index ad95ff3..481df58 100644 (file)
--- a/p2pool/bitcoin/data.py
+++ b/p2pool/bitcoin/data.py
@@ -12,22 +12,24 @@ def hash160(data):
     return pack.IntType(160).unpack(hashlib.new('ripemd160', hashlib.sha256(data).digest()).digest())
 
 class ChecksummedType(pack.Type):
-    def __init__(self, inner):
+    def __init__(self, inner, checksum_func=lambda data: hashlib.sha256(hashlib.sha256(data).digest()).digest()[:4]):
         self.inner = inner
+        self.checksum_func = checksum_func
     
     def read(self, file):
         obj, file = self.inner.read(file)
         data = self.inner.pack(obj)
         
-        checksum, file = pack.read(file, 4)
-        if checksum != hashlib.sha256(hashlib.sha256(data).digest()).digest()[:4]:
+        calculated_checksum = self.checksum_func(data)
+        checksum, file = pack.read(file, len(calculated_checksum))
+        if checksum != calculated_checksum:
             raise ValueError('invalid checksum')
         
         return obj, file
     
     def write(self, file, item):
         data = self.inner.pack(item)
-        return (file, data), hashlib.sha256(hashlib.sha256(data).digest()).digest()[:4]
+        return (file, data), self.checksum_func(data)
 
 class FloatingInteger(object):
     __slots__ = ['bits', '_target']