File: extend_json.py

package info (click to toggle)
python-bitarray 3.6.1-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 1,288 kB
  • sloc: python: 11,456; ansic: 7,657; makefile: 73; sh: 6
file content (56 lines) | stat: -rw-r--r-- 1,438 bytes parent folder | download | duplicates (2)
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
import json
from base64 import b64encode, b64decode

from bitarray import bitarray
from bitarray.util import serialize, deserialize


class JSONEncoder(json.JSONEncoder):

    def default(self, obj):

        if isinstance(obj, bitarray):
            if len(obj) > 50:
                return {'bitarray_b64': b64encode(serialize(obj)).decode()}
            else:
                return {'bitarray': obj.to01()}

        return json.JSONEncoder.default(self, obj)


class JSONDecoder(json.JSONDecoder):

    def __init__(self, *args, **kwargs):
        json.JSONDecoder.__init__(self, object_hook=self.object_hook,
                                  *args, **kwargs)

    def object_hook(self, obj):
        if isinstance(obj, dict) and len(obj) == 1:
            if 'bitarray_b64' in obj:
                return deserialize(b64decode(obj['bitarray_b64']))

            if 'bitarray' in obj:
                return bitarray(obj['bitarray'])

        return obj


def test():
    from random import getrandbits
    from bitarray.util import urandom

    a = [urandom(n * n, endian=['little', 'big'][getrandbits(1)])
         for n in range(12)]
    a.append({'key1': bitarray('010'),
              'key2': 'value2',
              'key3': urandom(300)})
    j = JSONEncoder(indent=2).encode(a)
    print(j)

    b = JSONDecoder().decode(j)
    assert a == b
    assert b[-1]['key1'] == bitarray('010')


if __name__ == '__main__':
    test()