codekingpro/portable-devtools

1"""2An implementation of a bitwise prefix tree specially built for decoding3Huffman-coded content where we already know the Huffman table.4"""5from __future__ import annotations6 7 8class HuffmanEncoder:9    """10    Encodes a string according to the Huffman encoding table defined in the11    HPACK specification.12    """13 14    def __init__(self, huffman_code_list: list[int], huffman_code_list_lengths: list[int]) -> None:15        self.huffman_code_list = huffman_code_list16        self.huffman_code_list_lengths = huffman_code_list_lengths17 18    def encode(self, bytes_to_encode: bytes | None) -> bytes:19        """20        Given a string of bytes, encodes them according to the HPACK Huffman21        specification.22        """23        # If handed the empty string, just immediately return.24        if not bytes_to_encode:25            return b""26 27        final_num = 028        final_int_len = 029 30        # Turn each byte into its huffman code. These codes aren't necessarily31        # octet aligned, so keep track of how far through an octet we are. To32        # handle this cleanly, just use a single giant integer.33        for byte in bytes_to_encode:34            bin_int_len = self.huffman_code_list_lengths[byte]35            bin_int = self.huffman_code_list[byte] & (36                2 ** (bin_int_len + 1) - 137            )38            final_num <<= bin_int_len39            final_num |= bin_int40            final_int_len += bin_int_len41 42        # Pad out to an octet with ones.43        bits_to_be_padded = (8 - (final_int_len % 8)) % 844        final_num <<= bits_to_be_padded45        final_num |= (1 << bits_to_be_padded) - 146 47        # Convert the number to hex and strip off the leading '0x' and the48        # trailing 'L', if present.49        s = hex(final_num)[2:].rstrip("L")50 51        # If this is odd, prepend a zero.52        s = "0" + s if len(s) % 2 != 0 else s53 54        # This number should have twice as many digits as bytes. If not, we're55        # missing some leading zeroes. Work out how many bytes we want and how56        # many digits we have, then add the missing zero digits to the front.57        total_bytes = (final_int_len + bits_to_be_padded) // 858        expected_digits = total_bytes * 259 60        if len(s) != expected_digits:61            missing_digits = expected_digits - len(s)62            s = ("0" * missing_digits) + s63 64        return bytes.fromhex(s)65