Improve Huffman algorithm by providing a way to compress strings not in a file, and adding doctests

Author: Lecdi

compression/huffman.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import sys
+import doctest
 
 
 class Letter:
@@ -20,25 +21,71 @@ def __init__(self, freq: int, left: Letter | TreeNode, right: Letter | TreeNode)
         self.right: Letter | TreeNode = right
 
 
+def parse_string(string: str) -> list[Letter]:
+    """
+    Return a list of Letter objects storing frequency
+    >>> string_in1 = "goose"
+    >>> out1 = parse_string(string_in1)
+    >>> out1
+    [g:1, s:1, e:1, o:2]
+    >>> string_in2 = ""
+    >>> out2 = parse_string(string_in2)
+    >>> out2
+    []
+    >>> string_in3 = "abbcccd"
+    >>> out3 = parse_string(string_in3)
+    >>> out3
+    [a:1, d:1, b:2, c:3]
+    """
+    chars: dict[str, Letter] = {}
+    for char in string:
+        if char not in chars:
+            chars[char] = Letter(char, 1)
+        else:
+            chars[char].freq += 1
+    return sorted(chars.values(), key=lambda letter: letter.freq)
+
+
 def parse_file(file_path: str) -> list[Letter]:
     """
-    Read the file and build a dict of all letters and their
-    frequencies, then convert the dict into a list of Letters.
+    Read file and return a list of Letter objects storing frequency
+    >>> test_file_path_in1 = "text_data/text_original.txt"
+    >>> out1 = parse_file(test_file_path_in1)
+    >>> out1
+    [T:1, h:1, a:1, e:1, i:2, t:2, s:3,  :3, .:3]
     """
-    chars: dict[str, int] = {}
-    with open(file_path) as f:
-        while True:
-            c = f.read(1)
-            if not c:
-                break
-            chars[c] = chars[c] + 1 if c in chars else 1
-    return sorted((Letter(c, f) for c, f in chars.items()), key=lambda x: x.freq)
+    with open(file_path, "r", encoding="utf8") as file:
+        string = file.read()
+    return parse_string(string)
 
 
 def build_tree(letters: list[Letter]) -> Letter | TreeNode:
     """
-    Run through the list of Letters and build the min heap
-    for the Huffman Tree.
+    Build the min heap for the Huffman Tree; return root node
+    >>> letters_in1 = [Letter('g', 1), Letter('s', 1), Letter('e', 1), Letter('o', 2)]
+    >>> out1 = build_tree(letters_in1)
+    >>> out1.freq
+    5
+    >>> out1.left.freq
+    2
+    >>> out1.left.left
+    g:1
+    >>> out1.left.right
+    s:1
+    >>> out1.right.freq
+    3
+    >>> out1.right.left
+    e:1
+    >>> out1.right.right
+    o:2
+    >>> letters_in2 = [Letter('a', 1), Letter('b', 1)]
+    >>> out2 = build_tree(letters_in2)
+    >>> out2.freq
+    2
+    >>> out2.left
+    a:1
+    >>> out2.right
+    b:1
     """
     response: list[Letter | TreeNode] = list(letters)
     while len(response) > 1:
@@ -51,10 +98,30 @@ def build_tree(letters: list[Letter]) -> Letter | TreeNode:
     return response[0]
 
 
-def traverse_tree(root: Letter | TreeNode, bitstring: str) -> list[Letter]:
+def traverse_tree(root: Letter | TreeNode, bitstring: str = "") -> list[Letter]:
     """
     Recursively traverse the Huffman Tree to set each
     Letter's bitstring dictionary, and return the list of Letters
+    >>> root_in1 = build_tree(parse_string("goose"))
+    >>> out1 = traverse_tree(root_in1, "")
+    >>> out1
+    [g:1, s:1, e:1, o:2]
+    >>> out1[0].bitstring['g']
+    '00'
+    >>> out1[1].bitstring['s']
+    '01'
+    >>> out1[2].bitstring['e']
+    '10'
+    >>> out1[3].bitstring['o']
+    '11'
+    >>> root_in2 = build_tree(parse_file("text_data/text_original.txt"))
+    >>> out2 = traverse_tree(root_in2)
+    >>> out2
+    [.:3, i:2, t:2, T:1, h:1, a:1, e:1, s:3,  :3]
+    >>> out2[0].bitstring['.']
+    '00'
+    >>> out2[4].bitstring['h']
+    '1001'
     """
     if isinstance(root, Letter):
         root.bitstring[root.letter] = bitstring
@@ -66,27 +133,40 @@ def traverse_tree(root: Letter | TreeNode, bitstring: str) -> list[Letter]:
     return letters
 
 
-def huffman(file_path: str) -> None:
+def huffman_string(string: str, sep=" ") -> str:
     """
-    Parse the file, build the tree, then run through the file
-    again, using the letters dictionary to find and print out the
-    bitstring for each letter.
+    Return huffman coded string, with
+    letter bitstrings separated by sep parameter
+    >>> huffman_string("goose")
+    '00 11 11 01 10'
+    >>> huffman_string("This is a test...", "")
+    '1000100101011011101011011110101110111011110011000000'
     """
-    letters_list = parse_file(file_path)
+    letters_list = parse_string(string)
     root = build_tree(letters_list)
-    letters = {
-        k: v for letter in traverse_tree(root, "") for k, v in letter.bitstring.items()
+    letter_bitstrings = {
+        k: v for letter in traverse_tree(root) for k, v in letter.bitstring.items()
     }
-    print(f"Huffman Coding  of {file_path}: ")
-    with open(file_path) as f:
-        while True:
-            c = f.read(1)
-            if not c:
-                break
-            print(letters[c], end=" ")
-    print()
+    return sep.join(letter_bitstrings[char] for char in string)
+
+
+def huffman(file_path: str) -> None:
+    """
+    Parse the file, huffman code it and print the result
+    >>> huffman("text_data/text_original.txt")
+    Huffman Coding of text_data/text_original.txt:
+    1000 1001 010 110 111 010 110 111 1010 111 011 1011 110 011 00 00 00
+    """
+    with open(file_path, "r", encoding="utf8") as file:
+        string = file.read()
+    result = huffman_string(string, " ")
+    print(f"Huffman Coding of {file_path}:\n{result}")
 
 
 if __name__ == "__main__":
-    # pass the file path to the huffman function
-    huffman(sys.argv[1])
+    if len(sys.argv) < 2:
+        # if no file path given, test the module
+        doctest.testmod()
+    else:
+        # pass the file path to the huffman function
+        huffman(sys.argv[1])

compression/text_data/text_original.txt

@@ -0,0 +1 @@
+This is a test...