fixed 5cde217 causing assertion failures in debug mode

[p2pool.git] / p2pool / util / pack.py

diff --git a/p2pool/util/pack.py b/p2pool/util/pack.py
index 0c50fc2..b2fd6da 100644 (file)
--- a/p2pool/util/pack.py
+++ b/p2pool/util/pack.py
@@ -2,6 +2,7 @@ import binascii
 import struct
 
 import p2pool
+from p2pool.util import memoize
 
 class EarlyEnd(Exception):
     pass
@@ -37,12 +38,12 @@ class Type(object):
     def __ne__(self, other):
         return not (self == other)
     
-    def _unpack(self, data):
+    def _unpack(self, data, ignore_trailing=False):
         obj, (data2, pos) = self.read((data, 0))
         
         assert data2 is data
         
-        if pos != len(data):
+        if pos != len(data) and not ignore_trailing:
             raise LateEnd()
         
         return obj
@@ -58,11 +59,13 @@ class Type(object):
         return ''.join(res)
     
     
-    def unpack(self, data):
-        obj = self._unpack(data)
+    def unpack(self, data, ignore_trailing=False):
+        obj = self._unpack(data, ignore_trailing)
         
         if p2pool.DEBUG:
-            if self._pack(obj) != data:
+            packed = self._pack(obj)
+            good = data.startswith(packed) if ignore_trailing else data == packed
+            if not good:
                 raise AssertionError()
         
         return obj
@@ -156,19 +159,21 @@ class EnumType(Type):
 class ListType(Type):
     _inner_size = VarIntType()
     
-    def __init__(self, type):
+    def __init__(self, type, mul=1):
         self.type = type
+        self.mul = mul
     
     def read(self, file):
         length, file = self._inner_size.read(file)
-        res = []
+        length *= self.mul
+        res = [None]*length
         for i in xrange(length):
-            item, file = self.type.read(file)
-            res.append(item)
+            res[i], file = self.type.read(file)
         return res, file
     
     def write(self, file, item):
-        file = self._inner_size.write(file, len(item))
+        assert len(item) % self.mul == 0
+        file = self._inner_size.write(file, len(item)//self.mul)
         for subitem in item:
             file = self.type.write(file, subitem)
         return file
@@ -187,6 +192,7 @@ class StructType(Type):
     def write(self, file, item):
         return file, struct.pack(self.desc, item)
 
+@memoize.fast_memoize_multiple_args
 class IntType(Type):
     __slots__ = 'bytes step format_str max'.split(' ')