1# Copyright 2017 CodiLime
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import msgpack
16import six
17import sys
18
19from veles.data.bindata import BinData
20from veles.compatibility import pep487
21from veles.schema import nodeid
22from veles.compatibility.int_bytes import int_to_bytes, int_from_bytes
23from veles.util.bigint import bigint_encode, bigint_decode
24
25
26EXT_NODE_ID = 0
27EXT_BINDATA = 1
28EXT_BIGINT = 2
29
30
31class MsgpackWrapper(pep487.NewObject):
32    def __init__(self):
33        self.packer = msgpack.Packer(
34            use_bin_type=True, default=MsgpackWrapper.pack_obj)
35        self.unpacker = msgpack.Unpacker(
36            ext_hook=MsgpackWrapper.load_obj, max_buffer_size=sys.maxsize)
37
38    @classmethod
39    def pack_obj(cls, obj):
40        if isinstance(obj, nodeid.NodeID):
41            return msgpack.ExtType(EXT_NODE_ID, obj.bytes)
42        if isinstance(obj, BinData):
43            width = int_to_bytes(obj.width, 4, 'little')
44            return msgpack.ExtType(EXT_BINDATA, width + obj.raw_data)
45        if isinstance(obj, six.integer_types):
46            return msgpack.ExtType(EXT_BIGINT, bigint_encode(obj))
47        raise TypeError('Object of unknown type {}'.format(obj))
48
49    @classmethod
50    def load_obj(cls, code, data):
51        if code == EXT_NODE_ID:
52            return nodeid.NodeID(data)
53        elif code == EXT_BINDATA:
54            width = int_from_bytes(data[:4], 'little')
55            return BinData.from_raw_data(width, data[4:])
56        elif code == EXT_BIGINT:
57            return bigint_decode(data)
58        return msgpack.ExtType(code, data)
59