convert to mmap, performance is worse

Author: gesslerpd

src/msgpack_stream/_io.py

@@ -13,7 +13,7 @@ def pack(obj):
 
 def unpack(data):
     """Unpack data into object."""
-    with io.BytesIO(data) as stream:
-        obj = unpack_stream(stream)
-        excess_data = stream.read()
+    with io.BytesIO(data) as stream, stream.getbuffer() as mv:
+        obj = unpack_stream(mv)
+        excess_data = bytes(mv[unpack_stream.__offset__:])
     return obj, excess_data

src/msgpack_stream/_msgpack.py

@@ -11,7 +11,7 @@
     f64_b_t,
 )
 
-# deference for performance
+# dereference for performance
 u8_b_pack = u8_b_t.pack
 u16_b_pack = u16_b_t.pack
 u32_b_pack = u32_b_t.pack
@@ -34,6 +34,17 @@
 f32_b_unpack = f32_b_t.unpack
 f64_b_unpack = f64_b_t.unpack
 
+u8_s = u8_b_t.struct
+u16_s = u16_b_t.struct
+u32_s = u32_b_t.struct
+u64_s = u64_b_t.struct
+s8_s = s8_b_t.struct
+s16_s = s16_b_t.struct
+s32_s = s32_b_t.struct
+s64_s = s64_b_t.struct
+f32_s = f32_b_t.struct
+f64_s = f64_b_t.struct
+
 
 def pack_stream(stream, obj):
     if isinstance(obj, dict):
@@ -141,63 +152,96 @@ def pack_stream(stream, obj):
         raise TypeError("type not supported:", obj, type(obj))
 
 
-def unpack_stream(stream):
-    first_byte = u8_b_unpack(stream)
+def unpack_stream(stream, offset=0):
+    (first_byte,) = u8_s.unpack_from(stream, offset)
+    offset += u8_s.size
     if first_byte <= 0x7F:
         obj = first_byte
     elif first_byte <= 0x8F:
         map_length = first_byte & 0b1111
-        obj = {unpack_stream(stream): unpack_stream(stream) for _ in range(map_length)}
+        unpack_stream.__offset__ = offset
+        obj = {
+            unpack_stream(stream, unpack_stream.__offset__): unpack_stream(
+                stream, unpack_stream.__offset__
+            )
+            for _ in range(map_length)
+        }
+        offset = unpack_stream.__offset__
     elif first_byte <= 0x9F:
         array_length = first_byte & 0b1111
-        obj = [unpack_stream(stream) for _ in range(array_length)]
+        unpack_stream.__offset__ = offset
+        obj = [
+            unpack_stream(stream, unpack_stream.__offset__) for _ in range(array_length)
+        ]
+        offset = unpack_stream.__offset__
     elif first_byte <= 0xBF:
         str_length = first_byte & 0b11111
-        obj = stream.read(str_length).decode("utf-8")
+        _offset = offset + str_length
+        obj = bytes(stream[offset:_offset]).decode("utf-8")
+        offset = _offset
     elif 0xE0 <= first_byte and first_byte <= 0xFF:
-        stream.seek(-1, 1)
-        obj = s8_b_unpack(stream)
+        (obj,) = s8_s.unpack_from(stream, offset - s8_s.size)
+        # peek backwards, offset is good where it's at
     elif first_byte == 0xC0:
         obj = None
     elif first_byte == 0xC2:
         obj = False
     elif first_byte == 0xC3:
         obj = True
     elif first_byte == 0xC4:
-        bin_length = u8_b_unpack(stream)
-        obj = stream.read(bin_length)
+        (bin_length,) = u8_s.unpack_from(stream, offset)
+        offset += u8_s.size
+        _offset = offset + bin_length
+        obj = bytes(stream[offset:_offset])
+        offset = _offset
     elif first_byte == 0xC5:
-        bin_length = u16_b_unpack(stream)
-        obj = stream.read(bin_length)
+        (bin_length,) = u16_s.unpack_from(stream, offset)
+        offset += u16_s.size
+        _offset = offset + bin_length
+        obj = bytes(stream[offset:_offset])
+        offset = _offset
     elif first_byte == 0xC6:
-        bin_length = u32_b_unpack(stream)
-        obj = stream.read(bin_length)
+        (bin_length,) = u32_s.unpack_from(stream, offset)
+        offset += u32_s.size
+        _offset = offset + bin_length
+        obj = bytes(stream[offset:_offset])
+        offset = _offset
     elif first_byte == 0xC7:
         raise NotImplementedError
     elif first_byte == 0xC8:
         raise NotImplementedError
     elif first_byte == 0xC9:
         raise NotImplementedError
     elif first_byte == 0xCA:
-        obj = f32_b_unpack(stream)
+        (obj,) = f32_s.unpack_from(stream, offset)
+        offset += f32_s.size
     elif first_byte == 0xCB:
-        obj = f64_b_unpack(stream)
+        (obj,) = f64_s.unpack_from(stream, offset)
+        offset += f64_s.size
     elif first_byte == 0xCC:
-        obj = u8_b_unpack(stream)
+        (obj,) = u8_s.unpack_from(stream, offset)
+        offset += u8_s.size
     elif first_byte == 0xCD:
-        obj = u16_b_unpack(stream)
+        (obj,) = u16_s.unpack_from(stream, offset)
+        offset += u16_s.size
     elif first_byte == 0xCE:
-        obj = u32_b_unpack(stream)
+        (obj,) = u32_s.unpack_from(stream, offset)
+        offset += u32_s.size
     elif first_byte == 0xCF:
-        obj = u64_b_unpack(stream)
+        (obj,) = u64_s.unpack_from(stream, offset)
+        offset += u64_s.size
     elif first_byte == 0xD0:
-        obj = s8_b_unpack(stream)
+        (obj,) = s8_s.unpack_from(stream, offset)
+        offset += s8_s.size
     elif first_byte == 0xD1:
-        obj = s16_b_unpack(stream)
+        (obj,) = s16_s.unpack_from(stream, offset)
+        offset += s16_s.size
     elif first_byte == 0xD2:
-        obj = s32_b_unpack(stream)
+        (obj,) = s32_s.unpack_from(stream, offset)
+        offset += s32_s.size
     elif first_byte == 0xD3:
-        obj = s64_b_unpack(stream)
+        (obj,) = s64_s.unpack_from(stream, offset)
+        offset += s64_s.size
     elif first_byte == 0xD4:
         raise NotImplementedError
     elif first_byte == 0xD5:
@@ -209,27 +253,63 @@ def unpack_stream(stream):
     elif first_byte == 0xD8:
         raise NotImplementedError
     elif first_byte == 0xD9:
-        str_length = u8_b_unpack(stream)
-        obj = stream.read(str_length).decode("utf-8")
+        (str_length,) = u8_s.unpack_from(stream, offset)
+        offset += u8_s.size
+        _offset = offset + str_length
+        obj = bytes(stream[offset:_offset]).decode("utf-8")
+        offset = _offset
     elif first_byte == 0xDA:
-        str_length = u16_b_unpack(stream)
-        obj = stream.read(str_length).decode("utf-8")
+        (str_length,) = u16_s.unpack_from(stream, offset)
+        offset += u16_s.size
+        _offset = offset + str_length
+        obj = bytes(stream[offset:_offset]).decode("utf-8")
+        offset = _offset
     elif first_byte == 0xDB:
-        str_length = u32_b_unpack(stream)
-        obj = stream.read(str_length).decode("utf-8")
+        (str_length,) = u32_s.unpack_from(stream, offset)
+        offset += u32_s.size
+        _offset = offset + str_length
+        obj = bytes(stream[offset:_offset]).decode("utf-8")
+        offset = _offset
     elif first_byte == 0xDC:
-        array_length = u16_b_unpack(stream)
-        obj = [unpack_stream(stream) for _ in range(array_length)]
+        (array_length,) = u16_s.unpack_from(stream, offset)
+        offset += u16_s.size
+        unpack_stream.__offset__ = offset
+        obj = [
+            unpack_stream(stream, unpack_stream.__offset__) for _ in range(array_length)
+        ]
+        offset = unpack_stream.__offset__
     elif first_byte == 0xDD:
-        array_length = u32_b_unpack(stream)
-        obj = [unpack_stream(stream) for _ in range(array_length)]
+        (array_length,) = u32_s.unpack_from(stream, offset)
+        offset += u32_s.size
+        unpack_stream.__offset__ = offset
+        obj = [
+            unpack_stream(stream, unpack_stream.__offset__) for _ in range(array_length)
+        ]
+        offset = unpack_stream.__offset__
     elif first_byte == 0xDE:
-        map_length = u16_b_unpack(stream)
-        obj = {unpack_stream(stream): unpack_stream(stream) for _ in range(map_length)}
+        (map_length,) = u16_s.unpack_from(stream, offset)
+        offset += u16_s.size
+        unpack_stream.__offset__ = offset
+        obj = {
+            unpack_stream(stream, unpack_stream.__offset__): unpack_stream(
+                stream, offset=unpack_stream.__offset__
+            )
+            for _ in range(map_length)
+        }
+        offset = unpack_stream.__offset__
     elif first_byte == 0xDF:
-        map_length = u32_b_unpack(stream)
-        obj = {unpack_stream(stream): unpack_stream(stream) for _ in range(map_length)}
+        (map_length,) = u32_s.unpack_from(stream, offset)
+        offset += u32_s.size
+        unpack_stream.__offset__ = offset
+        obj = {
+            unpack_stream(stream, unpack_stream.__offset__): unpack_stream(
+                stream, offset=unpack_stream.__offset__
+            )
+            for _ in range(map_length)
+        }
+        offset = unpack_stream.__offset__
     else:
         raise ValueError("invalid first byte", first_byte, hex(first_byte))
 
+    unpack_stream.__offset__ = offset
     return obj

tests/test_stream.py

@@ -23,10 +23,9 @@
 
 
 def test_stream():
-    stream = io.BytesIO()
-    pack_stream(stream, obj)
-    pack_stream(stream, obj)
-    stream.seek(0)
-    assert unpack_stream(stream) == obj
-    assert unpack_stream(stream) == obj
-    stream.close()
+    with io.BytesIO() as stream:
+        pack_stream(stream, obj)
+        pack_stream(stream, obj)
+        with stream.getbuffer() as mv:
+            assert unpack_stream(mv) == obj
+            assert unpack_stream(mv) == obj