LZW Decompression Library

Question

I've been asked to write a decompressor library for an interview process. It's a LZW decompression library. I've got it working and tried it out. I'm looking for any feedback on structure, presentation, readability from professionals who've worked as software engineers.

LZW explanation. The second table in uncompression section explains how the dictionary is built and key size is limited to 12-bits in my case.

"""
Decompressor library for Lempel–Ziv–Welch compression algorithm.
Main Functions:
    decompress(file_path)
    extract_archive_to_file(input_path, output_path)
Test code in the end.
"""


def read_compressed_codes_from_archive(file_path: str) -> [int]:
    """
    Reads file at specified path and returns content as an int array.
    :param file_path: path to archive
    :return: array of int codes
    """
    compressed_codes = []
    with open(file_path, "rb") as archive:
        # Read the 3 bytes, which contain two 12-bit codes
        bytes = archive.read(3)
        while bytes:
            if len(bytes) == 3:
                code1, code2 = unpack_codes_from_3_bytes(bytes)

                compressed_codes.append(code1)
                compressed_codes.append(code2)
            elif len(bytes) == 2:
                # Odd number of codes in the end of the file
                # Extract a single 12-bit code which is padded to 16-bits
                code = (0b0000111111111111 & (bytes[0] << 8)) | bytes[1]
                compressed_codes.append(code)

            bytes = archive.read(3)

    return compressed_codes


def unpack_codes_from_3_bytes(packed_bytes: [bytes]) -> (int, int):
    """
    Extracts two 12-bit numbers from 3 bytes (24-bits)
    :param packed_bytes: array of 3 int bytes
    :return: extracted int codes
    """
    # Create 16-bit int vars to store extracted codes in
    code1 = 0b0000111111111111
    code2 = 0b0000111111111111

    # Copy over 8 most significant bits and shift them to correct place
    code1 = code1 & packed_bytes[0] << 4
    # Copy over the 4 least significant bits from second byte
    code1 = code1 | (packed_bytes[1] >> 4)

    # Copy over the 4 most significant bits from second byte
    code2 = code2 & packed_bytes[1] << 8
    # Copy over the 8 least significant bits from third byte
    code2 = code2 | packed_bytes[2]

    return code1, code2


def convert_codes_to_text(compressed_codes: [int]) -> str:
    """
    Decompresses codes into text data.
    :param compressed_codes: compressed codes from archive
    :return: string data from compressed codes
    """
    # Create initial dictionary with individual char compressed_codes
    dictionary_size = 256
    dictionary = {i: chr(i) for i in range(dictionary_size)}

    text_data = ""
    previous_code = None

    # Loop over compressed compressed_codes and convert them to text values.
    for code in compressed_codes:
        if code in dictionary:
            text_data += dictionary[code]

            if previous_code:
                # Make sure new value is not already in the dict,
                # otherwise add it to the end.
                if dictionary[previous_code] + dictionary[code][0] not in \
                        dictionary.keys():
                    dictionary[dictionary_size] = dictionary[previous_code] \
                                                + dictionary[code][0]

                    dictionary_size += 1
                    # If dictionary is full, reset back to initial dict
                    # That is 4096 in binary ie one and 12 zeros
                    if dictionary_size == 0b1000000000000:
                        dictionary_size = 256
        else:
            assert ("Decompression error. Code {} not present in "
                    "dictionary".format(code))

        previous_code = code

    return text_data


def decompress(file_path: str) -> str:
    """
    Decompress an archive and get the text data.
    :param filepath: path to archive
    :return: text data as string
    """
    codes = read_compressed_codes_from_archive(file_path)
    text_data = convert_codes_to_text(codes)
    return text_data


def extract_archive_to_file(input_path: str, output_path: str):
    """
    Extracts an archive and saves it on disk.
    :param input_path: path to archive
    :param output_path: path to extracted file
    :return: none
    """
    text_data = decompress(input_path)
    with open(output_path, "w") as output_file:
        output_file.write(text_data)


def test_library():
    """
    Example test function.
    """
    file_path = "LzwInputData/compressedfile3.z"
    text_data = decompress(file_path)
    print(text_data)

    output_path = "LzwInputData/compressedfile3.txt"
    extract_archive_to_file(file_path, output_path)

test_library()

Thanks

Answer 1

Generally this looks like well-written, well documented, clean python code - so overall good job! But if I have to be picky, there are a few things I could mention...

code1, code2 = unpack_codes_from_3_bytes(bytes)
compressed_codes.append(code1)
compressed_codes.append(code2)

Could be shortened to:

compressed_codes.extend(unpack_codes_from_3_bytes(bytes))

I'm a bit uncertain as to why you do:

if dictionary_size == 0b1000000000000:

At the start, you define it as an int (dictionary_size = 256) then the only modification is int addition (dictionary_size += 1) so why the binary comparison?

Can you explain why you're using assert()? In a module like this I'd usually prefer to either raise exceptions explicitly - even going so far as to generate custom exception classes so that the end-user of the module can catch them easily - or if this is for reporting, use logging.error or similar.

There are few places where you generate variables and only use them once, e.g.:

text_data = convert_codes_to_text(codes)
return text_data

could be:

return convert_codes_to_text(codes)

or:

text_data = decompress(input_path)
with open(output_path, "w") as output_file:
    output_file.write(text_data)

could be:

with open(output_path, "w") as output_file:
    output_file.write(decompress(input_path))

Finally, not in dictionary.keys(): can just be not in dictionary:

Answer 2

First, I'd recommend that you change the API of read_compressed_codes_from_archive to not take the file-path, but a file-like object; see this post. It leaves you a lot more freedom if you someday want to get away from only reading from a file. As an example; an application (not in Python) that I used to work on eventually ran into issues where the 3rd-party API we used only accepted filenames, and we eventually got bug reports that it crashed for large files. It turns out that the way they had designed their library, they could only process the file by loading the entire thing at once, instead of over a stream or loading the file incrementally. Explicitly supporting a stream or something similar from the get-go is a good way to handle it.

Next, you should make the reader a generator, for very similar reasons. Overall, I'd probably rewrite the code to look something like this (note, I don't know what the correct doc and typing syntax for generators is, so I left as-is).

import typing

def read_compressed_codes_from_archive(archive: typing.BinaryIO) -> [int]:
    """
    Reads file at specified path and returns content as an int array.
    :param file_path: path to archive
    :return: array of int codes
    """

    bytes = archive.read(3)

    while bytes:
        if len(bytes) == 3:
            code1, code2 = unpack_codes_from_3_bytes(bytes)

            yield code1
            yield code2
        elif len(bytes) == 2:
            yield (TWELVE_BIT_MASK & (bytes[0] << 8)) | bytes[1]

        bytes = archive.read(3)

I also defined a global constant, TWELVE_BIT_MASK = 0b0000111111111111 so it's more clear what is actually happening there. Then, we can simplify the bit-twiddling a bit; getting the first 12-bit number doesn't require any masking. From there, I'd then move all of the bit twiddling into unpack_codes_from_3_bytes, and use yield from to get all the values, regardless of case. Finally, we can write a function to handle each use case of getting the first twelve and last twelve bits, like so:

import typing

def read_compressed_codes_from_archive(archive: typing.BinaryIO) -> [int]
    bytes = archive.read(3)

    while bytes:
        yield from unpack_codes_from_3_bytes(bytes)
        bytes = archive.read(3)

def unpack_codes_from_3_bytes(packed_bytes: [bytes]) -> int:

    if len(bytes) == 3:
        yield get_first_twelve_bits(packed_bytes)
    yield get_last_twelve_bits(packed_bytes)

def get_first_twelve_bits(packed_bytes: [bytes]) -> int:
    return (packed_bytes[0] << 4) | (packed_bytes[1] >> 4)

def get_last_twelve_bits(packed_bytes: [bytes]) -> int:
    return (TWELVE_BIT_MASK & packed_bytes[-2] << 8) | packed_bytes[-1]

You might even consider creating a class to represent an iterable LZW archive, like so (off the top of my head, so may not be correct)

class LzwArchive:

    def __init__(self, archive: typing.BinaryIO):
        self.archive = archive

    def __iter__(self):
        codes = self.archive.read(3)

        while codes:
            yield from LzwArchive._unpack_codes(codes)
            codes = self.archive.read(3)

    @staticmethod
    def _unpack_codes(codes: [bytes]) -> [int]:
        if len(codes) == 3:
            yield LzwArchive._get_first_twelve_bits(codes)
        yield LzwArchive._get_last_twelve_bits(codes)

    @staticmethod
    def _get_first_twelve_bits(packed_bytes: [bytes]) -> int:
        return (packed_bytes[0] << 4) | (packed_bytes[1] >> 4)

    @staticmethod
    def _get_last_twelve_bits(packed_bytes: [bytes]) -> int:
        return (TWELVE_BIT_MASK & packed_bytes[-2] << 8) | packed_bytes[-1]

Then you could just do something like this:

with open(myfile, 'rb') as f:
    archive = LzwArchive(f)
    for code in archive:
        process(code)

Let's keep the generators going! We can make convert_codes_to_text an iterable as well

def convert_codes_to_text(compressed_codes: [int]) -> str:
    dictionary_size = MIN_DICT_SIZE
    dictionary = {i: chr(i) for i in range(dictionary_size)}

    text_data = ""
    previous_code = None

    # Loop over compressed compressed_codes and convert them to text values.
    for code in compressed_codes:
        try:
            yield dictionary[code]
        except KeyError:
            # Do your thing here, but something better than an assert

        if previous_code:
            new_code = dictionary[code][0] + dictionary[previous_code]
            if new_code not in dictionary:
                dictionary[dictionary_size] = new_code
                dictionary_size += 1

            if dictionary_size = MAX_DICT_SIZE:
                dictionary_size = MIN_DICT_SIZE

        previous_code = code

It feels like we're doing a little too much work to maintain our dictionary here, so I think we can make a smarter data structure there as well. We could monkey around with LzwDictionary to make it subclass dict and truly quack like a dict (actually, subclassing defaultdict and/or OrderedDict can get you something pretty cool as well) that's more than you need right now.

def convert_codes_to_text(compressed_codes: [int]) -> str:
    dictionary = LzwDictionary()
    previous_code = None

    for code in compressed_codes:
        try:
            yield dictionary[code]
        except KeyError:
            # Do your thing here, but something better than an assert

        if previous_code:
            new_code = dictionary[code][0] + dictionary[previous_code]
            if new_code not in dictionary:
                dictionary.add_code(new_code)

        previous_code = code

class LzwDictionary:
    MAX_SIZE = 4096
    MIN_SIZE = 256

    def __init__(self):
        self.data = {i: chr(i) for i in range(LzwDictionary.MIN_SIZE)}
        self.size = LzwDictionary.MIN_SIZE

    def __getitem__(self, key):
        return self.data[key]

    def __contains__(self, key):
        return key in self.data

    def add_code(self, code):
        self.data[self.size] = code
        self.size += 1
        self._check_size()

    def _check_size(self):
        if self.size > LzwDictionary.MAX_SIZE:
            self.size = LzwDictionary.MIN_SIZE

At this point we just need to actually use the iterator; you could either do ''.join(convert_codes_to_text(codes)) or write incrementally. I'd probably end up with something like this:

def extract_archive_to_file(input_path, output_path):

    with open(input_path, 'rb') as lz_file:
        archive = LzwArchive(lz_file)
        decompressed = convert_codes_to_text(archive)
        with open(output_path, 'w') as out_file:
            for character in decompressed:
                out_file.write(character)

Wrote all of this without access to Python, so please forgive any typos or errors.