gif_lzw.py

class BitWriter():
    def __init__(self, consumer):
        self.consumer = consumer
        self.bit_buffer = 0
        self.bit_buffer_len = 0

    def write_bits(self, value, bits):
        # TODO catch invalid parameter
        self.bit_buffer |= value << self.bit_buffer_len
        self.bit_buffer_len += bits
        while self.bit_buffer_len >= 8:
            self.consumer.write(chr(self.bit_buffer & 0xff))
            self.bit_buffer >>= 8
            self.bit_buffer_len -= 8

    def finish(self):
        if self.bit_buffer_len > 0:
            self.consumer.write(chr(self.bit_buffer & 0xff))

def lzw_encode_default(data, code_bits):
    """
    @data: bytes (in py3) or something iterable with ints <= 255
    @code_bits: the minimum_lzw_size (size of code in bits at start = code_bits +1)
    data should contain the frames pixel index data.
    The result is a string (or bytes in py3).
    TODO: replace StringIO with BytesIo
    """
    MAX_DICT_SIZE = 4096
    buf = StringIO()
    writer = BitWriter(buf)

    if code_bits < 2 or code_bits > 8:
        raise Exception("Invalid code size");

    code_size = code_bits +1
    clear_code = 1 << code_bits
    eoi_code = clear_code +1
    next_code = eoi_code +1
    # We could use lists or arrays or something more memory friendly here i guess ?
    # Also we don't really need to recreate all the lists. Zeroing all fields should be enough ?
    root = list()
    for x in xrange(1 << code_size):
        root.append([None]*256)
    writer.write_bits(clear_code, code_size)

    i = 0
    while i < len(data):
#        print("now index #%i (%i)" % (i, data[i]))
        last_index = data[i]
        o = i + 1
        # If we are reading the last byte just write it out as it is.
        if o >= len(data):
            writer.write_bits(data[i], code_size)
            break

        # Loop till we find an unknown sequence
#        print("Do we know %i -> %i ?" % (last_index, data[o]))
        while (o < len(data) -1 and not root[last_index][ data[o] ] is None):
#            print("Known %i -> %i" % (last_index, root[last_index][ data[o] ]))
            last_index = root[last_index][ data[o] ]
            o += 1
#            print("Do we know %i -> %i ?" % (last_index, data[o]))

#        print("Create chain from %i to %i (%i)" % (last_index, data[o], next_code))
        root[last_index][data[o]] = next_code
        writer.write_bits(last_index, code_size)
        i = o

        #// Add new dictionary entry
        if next_code < MAX_DICT_SIZE:
            if (next_code & (next_code - 1)) == 0:  #// Is a power of 2
                code_size += 1
                # TODO: clean this up # Should be way nicer to calc.. to tired TODO:
                for x in xrange((1 << code_size) - (1 << (code_size-1))):
                    root.append([None]*256)
#                 print("increase code size to", self.codeBits)
            next_code += 1
        else:
            writer.write_bits(clear_code, code_size)
            code_size = code_bits +1
            next_code = eoi_code +1
            root = list()
            for x in xrange(1 << code_size):
                root.append([None]*256)

    writer.write_bits(eoi_code, code_size)
    writer.finish()
    return buf.getvalue()