1import array
2
3
4def packPiecefield(data):
5    if not isinstance(data, bytes) and not isinstance(data, bytearray):
6        raise Exception("Invalid data type: %s" % type(data))
7
8    res = []
9    if not data:
10        return array.array("H", b"")
11
12    if data[0] == b"\x00":
13        res.append(0)
14        find = b"\x01"
15    else:
16        find = b"\x00"
17    last_pos = 0
18    pos = 0
19    while 1:
20        pos = data.find(find, pos)
21        if find == b"\x00":
22            find = b"\x01"
23        else:
24            find = b"\x00"
25        if pos == -1:
26            res.append(len(data) - last_pos)
27            break
28        res.append(pos - last_pos)
29        last_pos = pos
30    return array.array("H", res)
31
32
33def unpackPiecefield(data):
34    if not data:
35        return b""
36
37    res = []
38    char = b"\x01"
39    for times in data:
40        if times > 10000:
41            return b""
42        res.append(char * times)
43        if char == b"\x01":
44            char = b"\x00"
45        else:
46            char = b"\x01"
47    return b"".join(res)
48
49
50def spliceBit(data, idx, bit):
51    if bit != b"\x00" and bit != b"\x01":
52        raise Exception("Invalid bit: %s" % bit)
53
54    if len(data) < idx:
55        data = data.ljust(idx + 1, b"\x00")
56    return data[:idx] + bit + data[idx+ 1:]
57
58class Piecefield(object):
59    def tostring(self):
60        return "".join(["1" if b else "0" for b in self.tobytes()])
61
62
63class BigfilePiecefield(Piecefield):
64    __slots__ = ["data"]
65
66    def __init__(self):
67        self.data = b""
68
69    def frombytes(self, s):
70        if not isinstance(s, bytes) and not isinstance(s, bytearray):
71            raise Exception("Invalid type: %s" % type(s))
72        self.data = s
73
74    def tobytes(self):
75        return self.data
76
77    def pack(self):
78        return packPiecefield(self.data).tobytes()
79
80    def unpack(self, s):
81        self.data = unpackPiecefield(array.array("H", s))
82
83    def __getitem__(self, key):
84        try:
85            return self.data[key]
86        except IndexError:
87            return False
88
89    def __setitem__(self, key, value):
90        self.data = spliceBit(self.data, key, value)
91
92class BigfilePiecefieldPacked(Piecefield):
93    __slots__ = ["data"]
94
95    def __init__(self):
96        self.data = b""
97
98    def frombytes(self, data):
99        if not isinstance(data, bytes) and not isinstance(data, bytearray):
100            raise Exception("Invalid type: %s" % type(data))
101        self.data = packPiecefield(data).tobytes()
102
103    def tobytes(self):
104        return unpackPiecefield(array.array("H", self.data))
105
106    def pack(self):
107        return array.array("H", self.data).tobytes()
108
109    def unpack(self, data):
110        self.data = data
111
112    def __getitem__(self, key):
113        try:
114            return self.tobytes()[key]
115        except IndexError:
116            return False
117
118    def __setitem__(self, key, value):
119        data = spliceBit(self.tobytes(), key, value)
120        self.frombytes(data)
121
122
123if __name__ == "__main__":
124    import os
125    import psutil
126    import time
127    testdata = b"\x01" * 100 + b"\x00" * 900 + b"\x01" * 4000 + b"\x00" * 4999 + b"\x01"
128    meminfo = psutil.Process(os.getpid()).memory_info
129
130    for storage in [BigfilePiecefieldPacked, BigfilePiecefield]:
131        print("-- Testing storage: %s --" % storage)
132        m = meminfo()[0]
133        s = time.time()
134        piecefields = {}
135        for i in range(10000):
136            piecefield = storage()
137            piecefield.frombytes(testdata[:i] + b"\x00" + testdata[i + 1:])
138            piecefields[i] = piecefield
139
140        print("Create x10000: +%sKB in %.3fs (len: %s)" % ((meminfo()[0] - m) / 1024, time.time() - s, len(piecefields[0].data)))
141
142        m = meminfo()[0]
143        s = time.time()
144        for piecefield in list(piecefields.values()):
145            val = piecefield[1000]
146
147        print("Query one x10000: +%sKB in %.3fs" % ((meminfo()[0] - m) / 1024, time.time() - s))
148
149        m = meminfo()[0]
150        s = time.time()
151        for piecefield in list(piecefields.values()):
152            piecefield[1000] = b"\x01"
153
154        print("Change one x10000: +%sKB in %.3fs" % ((meminfo()[0] - m) / 1024, time.time() - s))
155
156        m = meminfo()[0]
157        s = time.time()
158        for piecefield in list(piecefields.values()):
159            packed = piecefield.pack()
160
161        print("Pack x10000: +%sKB in %.3fs (len: %s)" % ((meminfo()[0] - m) / 1024, time.time() - s, len(packed)))
162
163        m = meminfo()[0]
164        s = time.time()
165        for piecefield in list(piecefields.values()):
166            piecefield.unpack(packed)
167
168        print("Unpack x10000: +%sKB in %.3fs (len: %s)" % ((meminfo()[0] - m) / 1024, time.time() - s, len(piecefields[0].data)))
169
170        piecefields = {}
171