Encode arbitrary strings with ANSI background coloring

color_code.py

#! /usr/bin/env python3
"""Encode arbitrary strings with ANSI background coloring."""

from argparse import ArgumentParser, Namespace
from sys import stdin, stdout


RESET = '\x1b[0m'


def color_code(data: bytes, *, end: str = RESET) -> str:
    """Color code the bytes."""

    return ''.join(
        f'\x1b[{int(octet) + 40}m ' for octet in reversed(
            oct(int.from_bytes(data, 'little'))[2:]
        )
    ) + end


def decode_colors(text: str) -> bytes:
    """Decode a color code."""

    return (i := int(''.join(
        map(lambda code: str(code - 40),
            filter(None, (
                int(code) for code in reversed(
                    text.rstrip(RESET)
                        .replace('\x1b[', '')
                        .replace('m', '')
                        .split()
                )
            ))
        )
    ), 8)).to_bytes(byte_length(i), 'little')


def byte_length(integer: int) -> int:
    """Returns the amount of necessary bytes to encode the int."""

    byte_count, remainder = divmod(integer.bit_length(), 8)
    return byte_count + bool(remainder)


def get_args(description: str = __doc__) -> Namespace:
    """Parse command line arguments."""

    parser = ArgumentParser(description=description)
    parser.add_argument('-d', '--decode', action='store_true')
    return parser.parse_args()


def main() -> None:
    """Run the script."""

    args = get_args()

    if args.decode:
        for line in stdin:
            stdout.buffer.write(decode_colors(line))
            stdout.flush()
    else:
        print(color_code(stdin.buffer.read()))


if __name__ == '__main__':
    main()