diff --git a/README.md b/README.md index 9f11bce..3cd9ee1 100644 --- a/README.md +++ b/README.md @@ -400,7 +400,7 @@ If a non-byte-string argument is passed to `umsgpack.unpackb()`, it will raise a * Python tuples and lists are both packed into the msgpack array format * Python float types are packed into the msgpack float32 or float64 format depending on the system's `sys.float_info` * The Python `datetime.datetime` type is packed into, and unpacked from, the msgpack `timestamp` format - * Note that this Python type only supports microsecond resolution, while the msgpack `timestamp` format supports nanosecond resolution. Timestamps with finer than microsecond resolution will lose precision during unpacking. + * Note that this Python type only supports microsecond resolution, while the msgpack `timestamp` format supports nanosecond resolution. Timestamps with finer than microsecond resolution will lose precision during unpacking. Users may override the packing and unpacking of the msgpack `timestamp` format with a custom type for alternate behavior. ## Testing diff --git a/test_umsgpack.py b/test_umsgpack.py index 926e0c6..eeea33e 100644 --- a/test_umsgpack.py +++ b/test_umsgpack.py @@ -334,6 +334,22 @@ b"\xd7\x30\x93\xc4\x03\x61\x62\x63\x7b\xc3"], ] +override_ext_handlers = { + datetime.datetime: + lambda obj: umsgpack.Ext(0x40, obj.strftime("%Y%m%dT%H:%M:%S.%f").encode()), + -0x01: + lambda ext: ext, +} + +override_ext_handlers_test_vectors = [ + ["pack override", + datetime.datetime(2000, 1, 1, 10, 5, 2, 0, umsgpack._utc_tzinfo), + b'\xc7\x18@20000101T10:05:02.000000'], + ["unpack override", + umsgpack.Ext(-0x01, b"\x00\xbb\xcc\xdd\x01\x02\x03\x04\x05\x06\x07\x08"), + b'\xc7\x0c\xff\x00\xbb\xcc\xdd\x01\x02\x03\x04\x05\x06\x07\x08'], +] + # These are the only global variables that should be exported by umsgpack exported_vars_test_vector = [ "Ext", @@ -492,10 +508,7 @@ def test_unpack_ordered_dict(self): def test_ext_exceptions(self): with self.assertRaises(TypeError): - _ = umsgpack.Ext(-1, b"") - - with self.assertRaises(TypeError): - _ = umsgpack.Ext(128, b"") + _ = umsgpack.Ext(5.0, b"") with self.assertRaises(TypeError): _ = umsgpack.Ext(0, u"unicode string") @@ -527,6 +540,26 @@ def test_pack_force_float_precision(self): packed = umsgpack.packb(obj, force_float_precision=precision) self.assertEqual(packed, data) + def test_pack_ext_override(self): + # Test overridden packing of datetime.datetime + (name, obj, data) = override_ext_handlers_test_vectors[0] + obj_repr = repr(obj) + print("\tTesting %s: object %s" % + (name, obj_repr if len(obj_repr) < 24 else obj_repr[0:24] + "...")) + + packed = umsgpack.packb(obj, ext_handlers=override_ext_handlers) + self.assertEqual(packed, data) + + def test_unpack_ext_override(self): + # Test overridden unpacking of Ext type -1 + (name, obj, data) = override_ext_handlers_test_vectors[1] + obj_repr = repr(obj) + print("\tTesting %s: object %s" % + (name, obj_repr if len(obj_repr) < 24 else obj_repr[0:24] + "...")) + + unpacked = umsgpack.unpackb(data, ext_handlers=override_ext_handlers) + self.assertEqual(unpacked, obj) + def test_streaming_writer(self): # Try first composite test vector (_, obj, data) = composite_test_vectors[0] diff --git a/umsgpack.py b/umsgpack.py index 7e2beb4..e7ff7a6 100644 --- a/umsgpack.py +++ b/umsgpack.py @@ -72,13 +72,9 @@ def __init__(self, type, data): Construct a new Ext object. Args: - type: application-defined type integer from 0 to 127 + type: application-defined type integer data: application-defined data byte array - Raises: - TypeError: - Specified ext type is outside of 0 to 127 range. - Example: >>> foo = umsgpack.Ext(0x05, b"\x01\x02\x03") >>> umsgpack.packb({u"special stuff": foo, u"awesome": True}) @@ -88,9 +84,9 @@ def __init__(self, type, data): Ext Object (Type: 0x05, Data: 01 02 03) >>> """ - # Application ext type should be 0 <= type <= 127 - if not isinstance(type, int) or not (type >= 0 and type <= 127): - raise TypeError("ext type out of range") + # Check type is type int + if not isinstance(type, int): + raise TypeError("ext type is not type integer") # Check data is type bytes elif sys.version_info[0] == 3 and not isinstance(data, bytes): raise TypeError("ext data is not type \'bytes\'") @@ -739,38 +735,38 @@ def _unpack_ext(code, fp, options): ext_type = struct.unpack("b", _read_except(fp, 1))[0] ext_data = _read_except(fp, length) - # Timestamp extension - if ext_type == -1: - return _unpack_ext_timestamp(code, ext_data, options) - - # Application extension + # Create extension object ext = Ext(ext_type, ext_data) # Unpack with ext handler, if we have one ext_handlers = options.get("ext_handlers") if ext_handlers and ext.type in ext_handlers: - ext = ext_handlers[ext.type](ext) + return ext_handlers[ext.type](ext) + + # Timestamp extension + if ext.type == -1: + return _unpack_ext_timestamp(ext, options) return ext -def _unpack_ext_timestamp(code, data, options): - if len(data) == 4: +def _unpack_ext_timestamp(ext, options): + if len(ext.data) == 4: # 32-bit timestamp - seconds = struct.unpack(">I", data)[0] + seconds = struct.unpack(">I", ext.data)[0] microseconds = 0 - elif len(data) == 8: + elif len(ext.data) == 8: # 64-bit timestamp - value = struct.unpack(">Q", data)[0] + value = struct.unpack(">Q", ext.data)[0] seconds = value & 0x3ffffffff microseconds = (value >> 34) // 1000 - elif len(data) == 12: + elif len(ext.data) == 12: # 96-bit timestamp - seconds = struct.unpack(">q", data[4:12])[0] - microseconds = struct.unpack(">I", data[0:4])[0] // 1000 + seconds = struct.unpack(">q", ext.data[4:12])[0] + microseconds = struct.unpack(">I", ext.data[0:4])[0] // 1000 else: raise UnsupportedTimestampException( - "unsupported timestamp with data length %d" % len(data)) + "unsupported timestamp with data length %d" % len(ext.data)) return _epoch + datetime.timedelta(seconds=seconds, microseconds=microseconds)