add ext support

Author: gesslerpd

src/msgpack_stream/__init__.py

@@ -2,3 +2,4 @@
 
 from ._io import pack as pack, unpack as unpack
 from ._msgpack import pack_stream as pack_stream, unpack_stream as unpack_stream
+from ._ext import ExtType as ExtType

src/msgpack_stream/_ext.py

@@ -0,0 +1,11 @@
+class ExtType:
+    __slots__ = ("code", "data")
+
+    def __init__(self, code: int, data: bytes):
+        self.code = code
+        self.data = data
+
+    def __eq__(self, other):
+        if isinstance(other, ExtType):
+            return self.code == other.code and self.data == other.data
+        return False

src/msgpack_stream/_msgpack.py

@@ -11,6 +11,7 @@
     f32_b_t,
     f64_b_t,
 )
+from ._ext import ExtType
 
 # deference for performance
 
@@ -49,6 +50,7 @@
 f64_b_unpack = f64_b_t.unpack
 
 _B = tuple([bytes([i]) for i in range(256)])
+_PO2 = {1: b"\xd4", 2: b"\xd5", 4: b"\xd6", 8: b"\xd7", 16: b"\xd8"}
 
 
 def pack_stream(stream, obj):
@@ -137,6 +139,21 @@ def pack_stream(stream, obj):
         else:
             raise ValueError("bin too large", obj)
         stream.write(obj)
+    elif _type is ExtType:
+        data = obj.data
+        p_code = s8_b_pack(obj.code)
+        extl = len(data)
+        if extl <= 16 and extl in _PO2:  # fixext (0xD4 - 0xD8)
+            stream.write(_PO2[extl] + p_code)
+        elif extl <= 0xFF:  # ext8
+            stream.write(b"\xc7" + _B[extl] + p_code)
+        elif extl <= 0xFF_FF:  # ext16
+            stream.write(b"\xc8" + u16_b_pack(extl) + p_code)
+        elif extl <= 0xFF_FF_FF_FF:  # ext32
+            stream.write(b"\xc9" + u32_b_pack(extl) + p_code)
+        else:
+            raise ValueError("ext too large", obj)
+        stream.write(data)
     else:
         raise TypeError("type not supported:", obj, _type)
 
@@ -172,8 +189,14 @@ def unpack_stream(stream):
     elif first_byte == 0xC6:
         bl = u32_b_unpack(stream)
         obj = stream.read(bl)
-    elif 0xC7 <= first_byte <= 0xC9:
-        raise NotImplementedError
+    elif first_byte <= 0xC9:  # ext (0xC7 - 0xC9)
+        if first_byte == 0xC7:  # ext8
+            extl = u8_b_unpack(stream)
+        elif first_byte == 0xC8:  # ext16
+            extl = u16_b_unpack(stream)
+        else:  # ext32
+            extl = u32_b_unpack(stream)
+        obj = ExtType(s8_b_unpack(stream), stream.read(extl))
     elif first_byte == 0xCA:
         obj = f32_b_unpack(stream)
     elif first_byte == 0xCB:
@@ -194,8 +217,8 @@ def unpack_stream(stream):
         obj = s32_b_unpack(stream)
     elif first_byte == 0xD3:
         obj = s64_b_unpack(stream)
-    elif 0xD4 <= first_byte <= 0xD8:
-        raise NotImplementedError
+    elif first_byte <= 0xD8:  # fixext (0xD4 - 0xD8)
+        obj = ExtType(s8_b_unpack(stream), stream.read(1 << (first_byte - 0xD4)))
     elif first_byte == 0xD9:
         sl = u8_b_unpack(stream)
         obj = stream.read(sl).decode("utf-8")

tests/test_ext.py

@@ -0,0 +1,83 @@
+import io
+import pytest
+from msgpack_stream import pack_stream, unpack_stream, ExtType
+
+
+def test_fixext1():
+    stream = io.BytesIO()
+    obj = ExtType(1, b"a")
+    pack_stream(stream, obj)
+    stream.seek(0)
+    assert stream.read(1) == b"\xd4"
+    stream.seek(0)
+    assert unpack_stream(stream) == obj
+
+
+def test_fixext2():
+    stream = io.BytesIO()
+    obj = ExtType(2, b"ab")
+    pack_stream(stream, obj)
+    stream.seek(0)
+    assert stream.read(1) == b"\xd5"
+    stream.seek(0)
+    assert unpack_stream(stream) == obj
+
+
+def test_fixext4():
+    stream = io.BytesIO()
+    obj = ExtType(3, b"abcd")
+    pack_stream(stream, obj)
+    stream.seek(0)
+    assert stream.read(1) == b"\xd6"
+    stream.seek(0)
+    assert unpack_stream(stream) == obj
+
+
+def test_fixext8():
+    stream = io.BytesIO()
+    obj = ExtType(4, b"a" * 8)
+    pack_stream(stream, obj)
+    stream.seek(0)
+    assert stream.read(1) == b"\xd7"
+    stream.seek(0)
+    assert unpack_stream(stream) == obj
+
+
+def test_fixext16():
+    stream = io.BytesIO()
+    obj = ExtType(5, b"a" * 16)
+    pack_stream(stream, obj)
+    stream.seek(0)
+    assert stream.read(1) == b"\xd8"
+    stream.seek(0)
+    assert unpack_stream(stream) == obj
+
+
+def test_ext8():
+    stream = io.BytesIO()
+    obj = ExtType(6, b"abc")
+    pack_stream(stream, obj)
+    stream.seek(0)
+    assert stream.read(1) == b"\xc7"
+    stream.seek(0)
+    assert unpack_stream(stream) == obj
+
+
+def test_ext16():
+    stream = io.BytesIO()
+    obj = ExtType(7, b"a" * 256)
+    pack_stream(stream, obj)
+    stream.seek(0)
+    assert stream.read(1) == b"\xc8"
+    stream.seek(0)
+    assert unpack_stream(stream) == obj
+
+
+def test_ext32():
+    stream = io.BytesIO()
+    obj = ExtType(8, b"a" * 65536)
+    pack_stream(stream, obj)
+    stream.seek(0)
+    assert stream.read(1) == b"\xc9"
+    stream.seek(0)
+    assert unpack_stream(stream) == obj