Speed improvements

kbinxml/bytebuffer.py

@@ -20,6 +20,11 @@ class ByteBuffer():
         else:
             return self.endian + str(count) + type
 
+    def get_bytes(self, count):
+        start = self.offset
+        self.offset += count
+        return self.data[start:self.offset]
+
     def get(self, type, count = None):
         ret = self.peek(type, count)
         size = calcsize(type)
@@ -33,6 +38,10 @@ class ByteBuffer():
         ret = unpack_from(fmt, self.data, self.offset)
         return ret[0] if count is None else ret
 
+    def append_bytes(self, data):
+        self.data.extend(data)
+        self.offset += len(data)
+
     def append(self, data, type, count = None):
         fmt = self._format_type(type, count)
         self.offset += calcsize(fmt)

kbinxml/kbinxml.py

@@ -124,13 +124,13 @@ class KBinXML():
 
     def data_grab_auto(self):
         size = self.dataBuf.get_s32()
-        ret = self.dataBuf.get('B', size)
+        ret = self.dataBuf.get_bytes(size)
         self.dataBuf.realign_reads()
         return ret
 
     def data_append_auto(self, data):
         self.dataBuf.append_s32(len(data))
-        self.dataBuf.append(data, 'B', len(data))
+        self.dataBuf.append_bytes(data)
         self.dataBuf.realign_writes()
 
     def data_grab_string(self):
@@ -309,7 +309,7 @@ class KBinXML():
                     name = unpack_sixbit(self.nodeBuf)
                 else:
                     length = (self.nodeBuf.get_u8() & ~64) + 1
-                    name = self.nodeBuf.get('B', length)
+                    name = self.nodeBuf.get_bytes(length)
                     name = bytes(name).decode(self.encoding)
                 debug_print(name)
 

kbinxml/sixbit.py

@@ -2,55 +2,27 @@
 from builtins import bytes
 from bitarray import bitarray
 
+charmap = '0123456789:ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz'
+bytemap = {charmap[i] : bytes(chr(i).encode()) for i in range(len(charmap))}
+
+@profile
 def pack_sixbit(string, byteBuf):
-    chars = str_to_sixbit(string)
+    chars = [bytemap[x] for x in string]
     bits = bitarray(endian='big')
     for c in chars:
-        bits.frombytes(c.encode())
+        bits.frombytes(c)
         # leave only the 6 bits we care for
         del bits[-8:-6]
     data = bytes(bits.tobytes())
-    byteBuf.append_u8(len(string))
-    byteBuf.append(data, 'B', len(data))
+    byteBuf.append_bytes((len(string),))
+    byteBuf.append_bytes(data)
 
 def unpack_sixbit(byteBuf):
     length = byteBuf.get_u8()
-    length_bytes = (length * 6 + 7) // 8
+    length_bits = length * 6
+    length_bytes = (length_bits + 7) // 8
     bitBuf = bitarray(endian='big')
-    bitBuf.frombytes(bytes(byteBuf.get('B', length_bytes)))
-    result = []
-    offset = 0
-    for i in range(length):
-        result.append(ord(bitBuf[offset:offset+6].tobytes()) >> (8 - 6))
-        offset += 6
-    return sixbit_to_str(result)
-
-# 0-9 for numbers, 10 is ':', 11 to 36 for capitals, 37 for underscore, 38-63 for lowercase
-def sixbit_to_str(decompressed):
-    string = ''
-    for d in decompressed:
-        if d <= 10:
-            d += ord('0')
-        elif d < 37:
-            d += 54
-        elif d == 37:
-            d += 58
-        else:
-            d += 59
-        string += chr(d)
-    return string
-
-def str_to_sixbit(string):
-    compress = []
-    for c in string:
-        if c >= '0' and c <= ':':
-            compress.append(ord(c) - ord('0'))
-        elif c >= 'A' and c <= 'Z':
-            compress.append(ord(c) - 54)
-        elif c == '_':
-            compress.append(ord(c) - 58)
-        elif c >= 'a' and c <= 'z':
-            compress.append(ord(c) - 59)
-        else:
-            raise ValueError('Node or attribute name can only contain alphanumeric + underscore')
-    return ''.join(map(chr, compress))
+    bitBuf.frombytes(bytes(byteBuf.get_bytes(length_bytes)))
+    result = [bitBuf[offset:offset+6].tobytes()[0] >> 2
+              for offset in range(0, length_bits, 6)]
+    return ''.join([charmap[x] for x in result])