avoid deep recursion in bdecode3, brings performance closer to py2 version

ox/torrent/bencode3.py

@@ -3,85 +3,93 @@
 # bencode.py python3 compatable bencode / bdecode
 #
 ##
+from collections import OrderedDict
 
-def _decode_int(data):
+
+class Decoder(object):
+
+    def _decode_int(self):
         """
         decode integer from bytearray
-    return int, remaining data
+        return int
         """
-    data = data[1:]
-    end = data.index(b'e')
-    return int(data[:end], 10), data[end+1:]
+        self.idx += 1
+        start = self.idx
+        end = self.data.index(b'e', self.idx)
+        self.idx = end + 1
+        return int(self.data[start:end])
 
-def _decode_str(data):
+    def _decode_str(self):
         """
         decode string from bytearray
-    return string, remaining data
+        return string
         """
-    start = data.index(b':')
-    l = int(data[:start].decode(), 10)
+        start = self.data.index(b':', self.idx)
+        l = int(self.data[self.idx:start].decode(), 10)
         if l <= 0:
             raise Exception('invalid string size: %d' % l)
         start += 1
-    ret = bytes(data[start:start+l])
+        ret = self.data[start:start+l]
         try:
             ret = ret.decode('utf-8')
         except:
             pass
-    data = data[start+l:]
-    return ret, data
+        self.idx = start + l
+        return ret
 
-def _decode_list(data):
+    def _decode_list(self):
         """
         decode list from bytearray
-    return list, remaining data
+        return list
         """
         ls = []
-    data = data[1:]
-    while data[0] != ord(b'e'):
-        elem, data = _decode(data)
-        ls.append(elem)
-    return ls, data[1:]
+        self.idx += 1
+        while self.data[self.idx] != ord(b'e'):
+            ls.append(self._decode())
+        self.idx += 1
+        return ls
 
-def _decode_dict(data):
+    def _decode_dict(self):
         """
         decode dict from bytearray
-    return dict, remaining data
+        return dict
         """
-    d = {}
-    data = data[1:]
-    while data[0] != ord(b'e'):
-        k, data = _decode_str(data)
-        v, data = _decode(data)
+        d = OrderedDict()
+        self.idx += 1
+        while self.data[self.idx] != ord(b'e'):
+            k = self._decode_str()
+            v = self._decode()
             d[k] = v
-    return d, data[1:]
+        self.idx += 1
+        return d
 
-def _decode(data):
-    """
-    decode a bytearray
-    return deserialized object, remaining data
-    """
-    ch = chr(data[0])
+    def _decode(self):
+        ch = chr(self.data[self.idx])
         if ch == 'l':
-        return _decode_list(data)
+            return self._decode_list()
         elif ch == 'i':
-        return _decode_int(data)
+            return self._decode_int()
         elif ch == 'd':
-        return _decode_dict(data)
+            return self._decode_dict()
         elif ch.isdigit():
-        return _decode_str(data)
+            return self._decode_str()
         else:
-        raise Exception('could not deserialize data: %s' % data)
+            raise Exception('could not decode data: %s' % data)
+
+    def decode(self, data):
+        self.idx = 0
+        self.data = data
+        obj = self._decode()
+        if len(data) != self.idx:
+            raise Exception('failed to decode, extra data: %s' % data)
+        return obj
 
 def bdecode(data):
     """
     decode a bytearray
-    return deserialized object
+    return decoded object
     """
-    obj, data = _decode(data)
-    if len(data) > 0:
-        raise Exception('failed to deserialize, extra data: %s' % data)
-    return obj
+    return Decoder().decode(data)
 
 def _encode_str(s, buff):
     """
@@ -114,10 +122,10 @@ def _encode_dict(d, buff):
     encode dict
     """
     buff.append(b'd')
-    l = list(d.keys())
-    l.sort()
-    for k in l:
-        _encode(str(k), buff)
+    for k in sorted(d):
+        if not isinstance(k, [bytes, str]):
+            k = str(k)
+        _encode(k, buff)
         _encode(d[k], buff)
     buff.append(b'e')