1*e1fe3e4aSElliott Hughesimport struct, warnings 2*e1fe3e4aSElliott Hughes 3*e1fe3e4aSElliott Hughestry: 4*e1fe3e4aSElliott Hughes import lz4 5*e1fe3e4aSElliott Hughesexcept ImportError: 6*e1fe3e4aSElliott Hughes lz4 = None 7*e1fe3e4aSElliott Hugheselse: 8*e1fe3e4aSElliott Hughes import lz4.block 9*e1fe3e4aSElliott Hughes 10*e1fe3e4aSElliott Hughes# old scheme for VERSION < 0.9 otherwise use lz4.block 11*e1fe3e4aSElliott Hughes 12*e1fe3e4aSElliott Hughes 13*e1fe3e4aSElliott Hughesdef decompress(data): 14*e1fe3e4aSElliott Hughes (compression,) = struct.unpack(">L", data[4:8]) 15*e1fe3e4aSElliott Hughes scheme = compression >> 27 16*e1fe3e4aSElliott Hughes size = compression & 0x07FFFFFF 17*e1fe3e4aSElliott Hughes if scheme == 0: 18*e1fe3e4aSElliott Hughes pass 19*e1fe3e4aSElliott Hughes elif scheme == 1 and lz4: 20*e1fe3e4aSElliott Hughes res = lz4.block.decompress(struct.pack("<L", size) + data[8:]) 21*e1fe3e4aSElliott Hughes if len(res) != size: 22*e1fe3e4aSElliott Hughes warnings.warn("Table decompression failed.") 23*e1fe3e4aSElliott Hughes else: 24*e1fe3e4aSElliott Hughes data = res 25*e1fe3e4aSElliott Hughes else: 26*e1fe3e4aSElliott Hughes warnings.warn("Table is compressed with an unsupported compression scheme") 27*e1fe3e4aSElliott Hughes return (data, scheme) 28*e1fe3e4aSElliott Hughes 29*e1fe3e4aSElliott Hughes 30*e1fe3e4aSElliott Hughesdef compress(scheme, data): 31*e1fe3e4aSElliott Hughes hdr = data[:4] + struct.pack(">L", (scheme << 27) + (len(data) & 0x07FFFFFF)) 32*e1fe3e4aSElliott Hughes if scheme == 0: 33*e1fe3e4aSElliott Hughes return data 34*e1fe3e4aSElliott Hughes elif scheme == 1 and lz4: 35*e1fe3e4aSElliott Hughes res = lz4.block.compress( 36*e1fe3e4aSElliott Hughes data, mode="high_compression", compression=16, store_size=False 37*e1fe3e4aSElliott Hughes ) 38*e1fe3e4aSElliott Hughes return hdr + res 39*e1fe3e4aSElliott Hughes else: 40*e1fe3e4aSElliott Hughes warnings.warn("Table failed to compress by unsupported compression scheme") 41*e1fe3e4aSElliott Hughes return data 42*e1fe3e4aSElliott Hughes 43*e1fe3e4aSElliott Hughes 44*e1fe3e4aSElliott Hughesdef _entries(attrs, sameval): 45*e1fe3e4aSElliott Hughes ak = 0 46*e1fe3e4aSElliott Hughes vals = [] 47*e1fe3e4aSElliott Hughes lastv = 0 48*e1fe3e4aSElliott Hughes for k, v in attrs: 49*e1fe3e4aSElliott Hughes if len(vals) and (k != ak + 1 or (sameval and v != lastv)): 50*e1fe3e4aSElliott Hughes yield (ak - len(vals) + 1, len(vals), vals) 51*e1fe3e4aSElliott Hughes vals = [] 52*e1fe3e4aSElliott Hughes ak = k 53*e1fe3e4aSElliott Hughes vals.append(v) 54*e1fe3e4aSElliott Hughes lastv = v 55*e1fe3e4aSElliott Hughes yield (ak - len(vals) + 1, len(vals), vals) 56*e1fe3e4aSElliott Hughes 57*e1fe3e4aSElliott Hughes 58*e1fe3e4aSElliott Hughesdef entries(attributes, sameval=False): 59*e1fe3e4aSElliott Hughes g = _entries(sorted(attributes.items(), key=lambda x: int(x[0])), sameval) 60*e1fe3e4aSElliott Hughes return g 61*e1fe3e4aSElliott Hughes 62*e1fe3e4aSElliott Hughes 63*e1fe3e4aSElliott Hughesdef bininfo(num, size=1): 64*e1fe3e4aSElliott Hughes if num == 0: 65*e1fe3e4aSElliott Hughes return struct.pack(">4H", 0, 0, 0, 0) 66*e1fe3e4aSElliott Hughes srange = 1 67*e1fe3e4aSElliott Hughes select = 0 68*e1fe3e4aSElliott Hughes while srange <= num: 69*e1fe3e4aSElliott Hughes srange *= 2 70*e1fe3e4aSElliott Hughes select += 1 71*e1fe3e4aSElliott Hughes select -= 1 72*e1fe3e4aSElliott Hughes srange //= 2 73*e1fe3e4aSElliott Hughes srange *= size 74*e1fe3e4aSElliott Hughes shift = num * size - srange 75*e1fe3e4aSElliott Hughes return struct.pack(">4H", num, srange, select, shift) 76*e1fe3e4aSElliott Hughes 77*e1fe3e4aSElliott Hughes 78*e1fe3e4aSElliott Hughesdef num2tag(n): 79*e1fe3e4aSElliott Hughes if n < 0x200000: 80*e1fe3e4aSElliott Hughes return str(n) 81*e1fe3e4aSElliott Hughes else: 82*e1fe3e4aSElliott Hughes return ( 83*e1fe3e4aSElliott Hughes struct.unpack("4s", struct.pack(">L", n))[0].replace(b"\000", b"").decode() 84*e1fe3e4aSElliott Hughes ) 85*e1fe3e4aSElliott Hughes 86*e1fe3e4aSElliott Hughes 87*e1fe3e4aSElliott Hughesdef tag2num(n): 88*e1fe3e4aSElliott Hughes try: 89*e1fe3e4aSElliott Hughes return int(n) 90*e1fe3e4aSElliott Hughes except ValueError: 91*e1fe3e4aSElliott Hughes n = (n + " ")[:4] 92*e1fe3e4aSElliott Hughes return struct.unpack(">L", n.encode("ascii"))[0] 93