Add compression

huffman.py

@@ -171,6 +171,50 @@ def __init__(self) -> None:
             if k == HUFFMAN_LUTBITS:
                 self.decoded_lut[i] = node
 
+    def compress(self, data: bytes) -> bytes:
+        """
+        Compresses data using the teeworlds frequency table
+        """
+        src_index = 0
+        end = len(data)
+        bits = 0
+        bitcount = 0
+        dst = bytearray()
+
+        if len(data) == 0:
+            return b''
+
+        symbol = data[src_index]
+        src_index += 1
+
+        while src_index < end:
+            bits |= self.nodes[symbol].bits << bitcount
+            bitcount += self.nodes[symbol].num_bits
+
+            symbol = data[src_index]
+            src_index += 1
+
+            while bitcount >= 8:
+                dst.append(bits & 0xFF)
+                bits >>= 8
+                bitcount -= 8
+
+        bits |= self.nodes[symbol].bits << bitcount
+        bitcount += self.nodes[symbol].num_bits
+        while bitcount >= 8:
+            dst.append(bits & 0xFF)
+            bits >>= 8
+            bitcount -= 8
+
+        bits |= self.nodes[HUFFMAN_EOF_SYMBOL].bits << bitcount
+        bitcount += self.nodes[HUFFMAN_EOF_SYMBOL].num_bits
+        while bitcount >= 8:
+            dst.append(bits & 0xFF)
+            bits >>= 8
+            bitcount -= 8
+        dst.append(bits)
+        return dst
+
     def decompress(self, data: bytes) -> bytes:
         """
         Decompresses data using the teeworlds frequency table

tests/basic_test.py

@@ -14,6 +14,12 @@ def test_huffman_hello_world():
     expected = b'hello world'
     assert decompressed == expected
 
+def test_huffman_hello_world_compress():
+    huffman = Huffman()
+    compressed = huffman.compress(b'hello world')
+    expected = bytes([174, 149, 19, 92, 9, 87, 194, 22, 177, 86, 220, 218, 34, 56, 185, 18, 156, 168, 184, 1])
+    assert compressed == expected
+
 def test_huffman_A():
     huffman = Huffman()
     compressed = bytes([188, 21, 55, 0])