Packing ExtType and some cleanup

2 files changed, 34 insertions, 54 deletions

diff --git a/test/test_extension.py b/test/test_extension.py
index f2fa363..2f85ce3 100644
--- a/test/test_extension.py
+++ b/test/test_extension.py
@@ -1,5 +1,7 @@
+from __future__ import print_function
 import array
 import msgpack
+from msgpack import ExtType
 
 
 def test_pack_ext_type():
@@ -11,54 +13,45 @@ def test_pack_ext_type():
     assert p(b'AB')       == b'\xd5\x42AB'         # fixext 2
     assert p(b'ABCD')     == b'\xd6\x42ABCD'       # fixext 4
     assert p(b'ABCDEFGH') == b'\xd7\x42ABCDEFGH'   # fixext 8
-    assert p(b'A'*16)     == b'\xd8\x42' + 'A'*16  # fixext 16
+    assert p(b'A'*16)     == b'\xd8\x42' + b'A'*16 # fixext 16
     assert p(b'ABC')      == b'\xc7\x03\x42ABC'        # ext 8
     assert p(b'A'*0x0123)     == b'\xc8\x01\x23\x42' + b'A'*0x0123 # ext 16
     assert p(b'A'*0x00012345) == b'\xc9\x00\x01\x23\x45\x42' + b'A'*0x00012345 # ext 32
 
 
-def test_unpack_extended_type():
-    class MyUnpacker(msgpack.Unpacker):
-        def read_extended_type(self, typecode, data):
-            return (typecode, data)
+def test_unpack_ext_type():
+    def check(b, expected):
+        assert msgpack.unpackb(b) == expected
 
-    def u(s):
-        unpacker = MyUnpacker()
-        unpacker.feed(s)
-        return unpacker.unpack_one()
-
-    assert u('\xd4\x42A')         == (0x42, 'A')        # fixext 1
-    assert u('\xd5\x42AB')        == (0x42, 'AB')       # fixext 2
-    assert u('\xd6\x42ABCD')      == (0x42, 'ABCD')     # fixext 4
-    assert u('\xd7\x42ABCDEFGH')  == (0x42, 'ABCDEFGH') # fixext 8
-    assert u('\xd8\x42' + 'A'*16) == (0x42, 'A'*16)     # fixext 16
-    assert u('\xc7\x03\x42ABC')   == (0x42, 'ABC')      # ext 8
-    assert (u('\xc8\x01\x23\x42' + 'A'*0x0123) ==
-            (0x42, 'A'*0x0123))                         # ext 16
-    assert (u('\xc9\x00\x01\x23\x45\x42' + 'A'*0x00012345) ==
-            (0x42, 'A'*0x00012345))                     # ext 32
+    check(b'\xd4\x42A',         ExtType(0x42, b'A'))        # fixext 1
+    check(b'\xd5\x42AB',        ExtType(0x42, b'AB'))       # fixext 2
+    check(b'\xd6\x42ABCD',      ExtType(0x42, b'ABCD'))     # fixext 4
+    check(b'\xd7\x42ABCDEFGH',  ExtType(0x42, b'ABCDEFGH')) # fixext 8
+    check(b'\xd8\x42' + b'A'*16, ExtType(0x42, b'A'*16))    # fixext 16
+    check(b'\xc7\x03\x42ABC',   ExtType(0x42, b'ABC'))      # ext 8
+    check(b'\xc8\x01\x23\x42' + b'A'*0x0123,
+          ExtType(0x42, b'A'*0x0123))                        # ext 16
+    check(b'\xc9\x00\x01\x23\x45\x42' + b'A'*0x00012345,
+          ExtType(0x42, b'A'*0x00012345))                   # ext 32
 
 
 def test_extension_type():
-    class MyPacker(msgpack.Packer):
-        def handle_unknown_type(self, obj):
-            if isinstance(obj, array.array):
-                typecode = 123 # application specific typecode
-                data = obj.tostring()
-                self.pack_ext_type(typecode, data)
-                return True
-
-    class MyUnpacker(msgpack.Unpacker):
-        def read_extended_type(self, typecode, data):
-            assert typecode == 123
-            obj = array.array('d')
-            obj.fromstring(data)
-            return obj
-
-    obj = [42, 'hello', array.array('d', [1.1, 2.2, 3.3])]
-    packer = MyPacker()
-    unpacker = MyUnpacker(None)
-    s = packer.pack(obj)
-    unpacker.feed(s)
-    obj2 = unpacker.unpack_one()
+    def default(obj):
+        print('default called', obj)
+        if isinstance(obj, array.array):
+            typecode = 123 # application specific typecode
+            data = obj.tostring()
+            return ExtType(typecode, data)
+        raise TypeError("Unknwon type object %r" % (obj,))
+
+    def ext_hook(code, data):
+        print('ext_hook called', code, data)
+        assert code == 123
+        obj = array.array('d')
+        obj.fromstring(data)
+        return obj
+
+    obj = [42, b'hello', array.array('d', [1.1, 2.2, 3.3])]
+    s = msgpack.packb(obj, default=default)
+    obj2 = msgpack.unpackb(s, ext_hook=ext_hook)
     assert obj == obj2
diff --git a/test/test_sequnpack.py b/test/test_sequnpack.py
index abc447a..af66b78 100644
--- a/test/test_sequnpack.py
+++ b/test/test_sequnpack.py
@@ -85,16 +85,3 @@ def test_readbytes():
     assert unpacker.read_bytes(3) == b'oob'
     assert unpacker.unpack() == ord(b'a')
     assert unpacker.unpack() == ord(b'r')
-
-def test_unpack_one():
-    unpacker = Unpacker()
-    unpacker.feed('\xda\x00\x03abc')
-    assert unpacker.unpack_one() == 'abc'
-    #
-    unpacker = Unpacker()
-    unpacker.feed('\xda\x00\x03abcd')
-    py.test.raises(ExtraData, "unpacker.unpack_one()")
-    #
-    unpacker = Unpacker()
-    unpacker.feed('\xda\x00\x03ab')
-    py.test.raises(UnpackValueError, "unpacker.unpack_one()")