You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

97 lines
2.7 KiB

4 years ago
  1. import io
  2. import logging
  3. logger = logging.getLogger(__name__)
  4. class CorruptDataError(Exception):
  5. pass
  6. class LZWDecoder:
  7. def __init__(self, fp):
  8. self.fp = fp
  9. self.buff = 0
  10. self.bpos = 8
  11. self.nbits = 9
  12. self.table = None
  13. self.prevbuf = None
  14. def readbits(self, bits):
  15. v = 0
  16. while 1:
  17. # the number of remaining bits we can get from the current buffer.
  18. r = 8-self.bpos
  19. if bits <= r:
  20. # |-----8-bits-----|
  21. # |-bpos-|-bits-| |
  22. # | |----r----|
  23. v = (v<<bits) | ((self.buff>>(r-bits)) & ((1<<bits)-1))
  24. self.bpos += bits
  25. break
  26. else:
  27. # |-----8-bits-----|
  28. # |-bpos-|---bits----...
  29. # | |----r----|
  30. v = (v<<r) | (self.buff & ((1<<r)-1))
  31. bits -= r
  32. x = self.fp.read(1)
  33. if not x: raise EOFError
  34. self.buff = ord(x)
  35. self.bpos = 0
  36. return v
  37. def feed(self, code):
  38. x = b''
  39. if code == 256:
  40. self.table = [bytes([i]) for i in range(256)] # 0-255
  41. self.table.append(None) # 256
  42. self.table.append(None) # 257
  43. self.prevbuf = b''
  44. self.nbits = 9
  45. elif code == 257:
  46. pass
  47. elif not self.prevbuf:
  48. try:
  49. x = self.prevbuf = self.table[code]
  50. except (TypeError, IndexError):
  51. # TypeError: table is None
  52. raise CorruptDataError()
  53. else:
  54. if code < len(self.table):
  55. x = self.table[code]
  56. self.table.append(self.prevbuf+x[:1])
  57. elif code == len(self.table):
  58. self.table.append(self.prevbuf+self.prevbuf[:1])
  59. x = self.table[code]
  60. else:
  61. raise CorruptDataError()
  62. l = len(self.table)
  63. if l == 511:
  64. self.nbits = 10
  65. elif l == 1023:
  66. self.nbits = 11
  67. elif l == 2047:
  68. self.nbits = 12
  69. self.prevbuf = x
  70. return x
  71. def run(self):
  72. while 1:
  73. try:
  74. code = self.readbits(self.nbits)
  75. except EOFError:
  76. break
  77. try:
  78. x = self.feed(code)
  79. except CorruptDataError:
  80. # just ignore corrupt data and stop yielding there
  81. break
  82. yield x
  83. logger.debug('nbits=%d, code=%d, output=%r, table=%r', self.nbits, code, x, self.table)
  84. # lzwdecode
  85. def lzwdecode(data):
  86. fp = io.BytesIO(data)
  87. return b''.join(LZWDecoder(fp).run())