Merge 9d2fb59 into 222130c

philiptzou · philiptzou · commit ea28bbdc72cd · 2015-05-19T13:13:44.000Z
diff --git a/wcf/records/base.py b/wcf/records/base.py
@@ -199,49 +199,73 @@ class ArrayRecord(Record):
         0xB1: ('UuidTextWithEndElement', 16, ''),
     }
 
-    def __init__(self, element, recordtype, data):
+    def __init__(self, element, data, attributes):
         self.element = element
-        self.recordtype = recordtype
         self.count = len(data)
         self.data = data
+        recordtype = None
+        for data in self.data:
+            if recordtype is None:
+                recordtype = data.type + 1
+            else:
+                assert recordtype == data.type + 1
+        self.recordtype = recordtype
+        self.attributes = []
 
     def to_bytes(self):
         """
+        >>> from wcf.records.text import Int32TextRecord
         >>> from wcf.records.elements import ShortElementRecord
-        >>> ArrayRecord(ShortElementRecord('item'), 0x8D, ['\\x01\\x00\\x00\\x00', '\\x02\\x00\\x00\\x00', '\\x03\\x00\\x00\\x00']).to_bytes()
+        >>> ArrayRecord(ShortElementRecord('item'), [Int32TextRecord(1), Int32TextRecord(2), Int32TextRecord(3)], []).to_bytes()
         b'\\x03@\\x04item\\x01\\x8d\\x03\\x01\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x03\\x00\\x00\\x00'
         """
         bt = super(ArrayRecord, self).to_bytes()
         bt += self.element.to_bytes()
+        for attrib in self.attributes:
+            bt += attrib.to_bytes()
         bt += EndElementRecord().to_bytes()
         bt += bytes(struct.pack(b'<B', self.recordtype))
         bt += MultiByteInt31(self.count).to_bytes()
         for data in self.data:
-            if isinstance(data, bytes):
-                bt += data
-            elif hasattr(data, 'to_bytes'):
-                bt += data.to_bytes()
-            else:
-                bt += data.encode('latin1')
-
+            bt += data.to_bytes()[1:]
         return bytes(bt)
 
     @classmethod
     def parse(cls, fp):
+        """
+        >>> from wcf.records import *
+        >>> from io import BytesIO
+        >>> buf = BytesIO(b'@\\x04item\\x01\\x8d\\x03\\x01\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x03\\x00\\x00\\x00')
+        >>> r = ArrayRecord.parse(buf)
+        >>> r
+        <ArrayRecord(type=0x3)>
+        >>> r.to_bytes()
+        b'\\x03@\\x04item\\x01\\x8d\\x03\\x01\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x03\\x00\\x00\\x00'
+        """
         element = struct.unpack(b'<B', fp.read(1))[0]
-        element = __records__[element].parse(fp)
+        element = Record.records[element].parse(fp)
+        attributes = []
+        while True:
+            type = struct.unpack(b'<B', fp.read(1))[0]
+            obj = Record.records[type].parse(fp)
+            if isinstance(obj, EndElementRecord):
+                break
+            elif isinstance(obj, Attribute):
+                attributes.append(obj)
+            else:
+                raise ValueError('unknown type: %s' % hex(type))
         recordtype = struct.unpack(b'<B', fp.read(1))[0]
         count = MultiByteInt31.parse(fp).value
         data = []
         for i in range(count):
-            data.append(__records__[recordtype-1].parse(fp))
-        return cls(element, recordtype, data)
+            data.append(Record.records[recordtype-1].parse(fp))
+        return cls(element, data, attributes)
 
     def __str__(self):
         """
         >>> from wcf.records.elements import ShortElementRecord
         >>> from wcf.records.text import Int32TextRecord
-        >>> str(ArrayRecord(ShortElementRecord('item'), 0x8D, [Int32TextRecord(1),Int32TextRecord(2),Int32TextRecord(3)]))
+        >>> str(ArrayRecord(ShortElementRecord('item'), [Int32TextRecord(1), Int32TextRecord(2), Int32TextRecord(3)], []))
         '<item>1</item><item>2</item><item>3</item>'
         """
         string = ''
diff --git a/wcf/records/text.py b/wcf/records/text.py
@@ -35,7 +35,6 @@
 import datetime
 import logging
 import uuid
-import sys
 
 try:
     from htmlentitydefs import codepoint2name
@@ -195,16 +194,10 @@ class UnicodeChars8TextRecord(Text):
     type = 0xB6
 
     def __init__(self, string):
-        if sys.version_info >= (3, 0, 0):
-            if isinstance(string, str):
-                self.value = string
-            else:
-                self.value = str(string)
+        if isinstance(string, str):
+            self.value = string
         else:
-            if isinstance(string, unicode):
-                self.value = string
-            else:
-                self.value = unicode(string)
+            self.value = str(string)
 
     def to_bytes(self):
         """
@@ -458,16 +451,10 @@ class Chars8TextRecord(Text):
     type = 0x98
 
     def __init__(self, value):
-        if sys.version_info >= (3, 0, 0):
-            if isinstance(value, str):
-                self.value = value
-            else:
-                self.value = str(value)
+        if isinstance(value, str):
+            self.value = value
         else:
-            if isinstance(value, unicode):
-                self.value = value
-            else:
-                self.value = unicode(value)
+            self.value = str(value)
 
     def __str__(self):
         return escape(self.value)
diff --git a/wcf/xml2records.py b/wcf/xml2records.py
@@ -48,7 +48,7 @@ def reset(self):
     def _parse_tag(self, tag):
         if ':' in tag:
             prefix = tag[:tag.find(':')]
-            name   = tag[tag.find(':')+1:]
+            name = tag[tag.find(':')+1:]
 
             if len(prefix) == 1:
                 cls_name = 'Element' + prefix.upper() + 'Record'
@@ -86,7 +86,7 @@ def _store_data(self, data, end=False):
         self.last_record.childs.append(textrecord)
         #if end:
         #    textrecord.type += 1
- 
+
     def _parse_data(self, data):
         data = data.strip()
         b64 = False
@@ -160,24 +160,24 @@ def _parse_data(self, data):
 
             base_diff = 62135596800.0
             dt = int((time.mktime(dt.timetuple()) - base) * 10 + ms)
-           
+
             return DateTimeTextRecord(dt, tz)
-        
+
         # text as fallback
         val = len(data)
         if val < 2**8:
-            return Chars8TextRecord(data)
+            return UnicodeChars8TextRecord(data)
         elif val < 2**16:
-            return Chars16TextRecord(data)
+            return UnicodeChars16TextRecord(data)
         elif val < 2**32:
-            return Chars32TextRecord(data)
+            return UnicodeChars32TextRecord(data)
 
     def _parse_attr(self, name, value):
 
         if ':' in name:
             prefix = name[:name.find(':')]
             name   = name[name.find(':')+1:]
-           
+
             if prefix == 'xmlns':
                 if value in inverted_dict:
                     return DictionaryXmlnsAttributeRecord(name,
@@ -215,25 +215,25 @@ def handle_starttag(self, tag, attrs):
         if self.data:
             self._store_data(self.data, False)
             self.data = None
-       
+
         el = self._parse_tag(tag)
         for n, v in attrs:
             el.attributes.append(self._parse_attr(n, v))
         self.last_record.childs.append(el)
         el.parent = self.last_record
         self.last_record = el
-   
+
     def handle_startendtag(self, tag, attrs):
         if self.data:
             self._store_data(self.data, False)
             self.data = None
-       
+
         el = self._parse_tag(tag)
         for n, v in attrs:
             el.attributes.append(self._parse_attr(n, v))
         self.last_record.childs.append(el)
         #self.last_record.childs.append(EndElementRecord())
-   
+
     def handle_endtag(self, tag):
         if self.data:
             self._store_data(self.data, True)
@@ -324,14 +324,14 @@ def parse(cls, data):
         else:
             raise ValueError("%s has an incompatible type %s" % (data,
                 type(data)))
-        
+
         p.feed(xml)
 
         return p.records
 
 if __name__ == '__main__':
     import sys
-   
+
     fp = sys.stdin
 
     if len(sys.argv) > 1:
