automatically find the best format to encode extended types

2 files changed, 38 insertions, 8 deletions

diff --git a/msgpack/fallback.py b/msgpack/fallback.py
index 101bd0f..f984dcd 100644
--- a/msgpack/fallback.py
+++ b/msgpack/fallback.py
@@ -545,13 +545,29 @@ class Packer(object):
         # overridden by subclasses
         return None
 
-    def pack_extended_type(self, fmt, typecode, data):
-        # for now we support only this. We should add support for the other
-        # fixext/ext formats
-        assert fmt == "ext 32"
+    def pack_extended_type(self, typecode, data):
         assert 0 <= typecode <= 127
-        N = len(data)
-        self._buffer.write(struct.pack('>BIB', 0xc9, N, typecode))
+        n = len(data)
+        if n == 1:
+            header = struct.pack(">BB", 0xd4, typecode) # fixext 1
+        elif n == 2:
+            header = struct.pack(">BB", 0xd5, typecode) # fixext 2
+        elif n == 4:
+            header = struct.pack(">BB", 0xd6, typecode) # fixext 4
+        elif n == 8:
+            header = struct.pack(">BB", 0xd7, typecode) # fixext 8
+        elif n == 16:
+            header = struct.pack(">BB", 0xd8, typecode) # fixext 16
+        elif n <= 2**8-1:
+            header = struct.pack(">BBB", 0xc7, n, typecode) # ext 8
+        elif n <= 2**16-1:
+            header = struct.pack(">BHB", 0xc8, n, typecode) # ext 16
+        elif n <= 2**32-1:
+            header = struct.pack(">BIB", 0xc9, n, typecode) # ext 32
+        else:
+            raise PackValueError("ext data too large")            
+        #
+        self._buffer.write(header)
         self._buffer.write(data)
 
     def pack(self, obj):
diff --git a/test/test_extension.py b/test/test_extension.py
index 0b26f8e..1908fa2 100644
--- a/test/test_extension.py
+++ b/test/test_extension.py
@@ -1,14 +1,28 @@
 import array
+import struct
 import msgpack
 
+def test_pack_extended_type():
+    def p(s):
+        packer = msgpack.Packer()
+        packer.pack_extended_type(0x42, s)
+        return packer._buffer.getvalue()
+    assert p('A')        == '\xd4\x42A'          # fixext 1
+    assert p('AB')       == '\xd5\x42AB'         # fixext 2
+    assert p('ABCD')     == '\xd6\x42ABCD'       # fixext 4
+    assert p('ABCDEFGH') == '\xd7\x42ABCDEFGH'   # fixext 8
+    assert p('A'*16)     == '\xd8\x42' + 'A'*16  # fixext 16
+    assert p('ABC')      == '\xc7\x03\x42ABC'        # ext 8
+    assert p('A'*0x0123)        == '\xc8\x01\x23\x42' + 'A'*0x0123 # ext 16
+    assert p('A'*0x00012345)    == '\xc9\x00\x01\x23\x45\x42' + 'A'*0x00012345 # ext 32
+
 def test_extension_type():
     class MyPacker(msgpack.Packer):
         def handle_unknown_type(self, obj):
             if isinstance(obj, array.array):
-                fmt = "ext 32"
                 typecode = 123 # application specific typecode
                 data = obj.tostring()
-                self.pack_extended_type(fmt, typecode, data)
+                self.pack_extended_type(typecode, data)
                 return True
 
     class MyUnpacker(msgpack.Unpacker):