grUtils.py 2.2 KB

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