forked from zarr-developers/numcodecs
-
Notifications
You must be signed in to change notification settings - Fork 0
/
msgpacks.py
83 lines (67 loc) · 2.51 KB
/
msgpacks.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import numpy as np
import msgpack
from .abc import Codec
from .compat import ensure_contiguous_ndarray
class MsgPack(Codec):
"""Codec to encode data as msgpacked bytes. Useful for encoding an array of Python
objects.
.. versionchanged:: 0.6
The encoding format has been changed to include the array shape in the encoded
data, which ensures that all object arrays can be correctly encoded and decoded.
Parameters
----------
use_single_float : bool, optional
Use single precision float type for float.
use_bin_type : bool, optional
Use bin type introduced in msgpack spec 2.0 for bytes. It also enables str8 type
for unicode.
raw : bool, optional
If true, unpack msgpack raw to Python bytes. Otherwise, unpack to Python str
by decoding with UTF-8 encoding.
Examples
--------
>>> import numcodecs
>>> import numpy as np
>>> x = np.array(['foo', 'bar', 'baz'], dtype='object')
>>> codec = numcodecs.MsgPack()
>>> codec.decode(codec.encode(x))
array(['foo', 'bar', 'baz'], dtype=object)
See Also
--------
numcodecs.pickles.Pickle, numcodecs.json.JSON, numcodecs.vlen.VLenUTF8
Notes
-----
Requires `msgpack <https://pypi.org/project/msgpack/>`_ to be installed.
"""
codec_id = 'msgpack2'
def __init__(self, use_single_float=False, use_bin_type=True, raw=False):
self.use_single_float = use_single_float
self.use_bin_type = use_bin_type
self.raw = raw
def encode(self, buf):
buf = np.asarray(buf)
items = buf.tolist()
items.append(buf.dtype.str)
items.append(buf.shape)
return msgpack.packb(items, use_bin_type=self.use_bin_type,
use_single_float=self.use_single_float)
def decode(self, buf, out=None):
buf = ensure_contiguous_ndarray(buf)
items = msgpack.unpackb(buf, raw=self.raw)
dec = np.empty(items[-1], dtype=items[-2])
dec[:] = items[:-2]
if out is not None:
np.copyto(out, dec)
return out
else:
return dec
def get_config(self):
return dict(id=self.codec_id,
raw=self.raw,
use_single_float=self.use_single_float,
use_bin_type=self.use_bin_type)
def __repr__(self):
return (
'MsgPack(raw={!r}, use_bin_type={!r}, use_single_float={!r})'
.format(self.raw, self.use_bin_type, self.use_single_float)
)