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 2762605..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(' ')
     
@@ -265,22 +271,26 @@ def get_record(fields):
                 if isinstance(other, dict):
                     return dict(self) == other
                 elif isinstance(other, _Record):
-                    return all(self[k] == other[k] for k in self.keys())
+                    for k in fields:
+                        if getattr(self, k) != getattr(other, k):
+                            return False
+                    return True
                 elif other is None:
                     return False
                 raise TypeError()
             def __ne__(self, other):
                 return not (self == other)
         _record_types[fields] = _Record
-    return _record_types[fields]()
+    return _record_types[fields]
 
 class ComposedType(Type):
     def __init__(self, fields):
-        self.fields = tuple(fields)
+        self.fields = list(fields)
         self.field_names = set(k for k, v in fields)
+        self.record_type = get_record(k for k, v in self.fields)
     
     def read(self, file):
-        item = get_record(k for k, v in self.fields)
+        item = self.record_type()
         for key, type_ in self.fields:
             item[key], file = type_.read(file)
         return item, file