-
Notifications
You must be signed in to change notification settings - Fork 0
/
packer.py
57 lines (41 loc) · 1.75 KB
/
packer.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
import msgpack
import msgpack_numpy
import cPickle as pickle
def packb(x, version, use_msgpack=True):
""" Pack an object x (that can contain numpy objects) """
if use_msgpack:
return msgpack.packb({'version': version, 'data': x}, default=msgpack_numpy.encode)
else:
return pickle.dumps({'version': version, 'data': x})
def fpackb(x, version, filepath, use_msgpack=True):
""" Pack an object x and save it to file (that can contain numpy objects) """
packed = packb(x, version, use_msgpack)
with open(filepath, 'w') as f:
f.write(packed)
def unpackb(packed, use_msgpack=True):
""" Unpack an object x (that can contain numpy objects) """
if use_msgpack:
dic = msgpack.unpackb(packed, object_hook=msgpack_numpy.decode)
else:
dic = pickle.load(packed)
return dic['version'], dic['data']
def unpackb_version(packed, expected_version, use_list=True, use_msgpack=True):
""" Unpack an object x (that can contain numpy objects) """
if use_msgpack:
dic = msgpack.unpackb(packed, object_hook=msgpack_numpy.decode, use_list=use_list)
else:
dic = pickle.load(packed)
package_version = dic['version']
if package_version != expected_version:
raise ValueError('Unexpected version ({0}) when reading package (expected: {1})'.format(package_version, expected_version))
return dic['data']
def funpackb_version(expected_version, filepath, use_list=True, use_msgpack=True):
""" Unpack an object x from file (that can contain numpy objects) """
with open(filepath, 'r') as f:
packed = f.read()
return unpackb_version(
packed=packed,
expected_version=expected_version,
use_list=use_list,
use_msgpack=use_msgpack,
)