like GNU `fmt` with ANSI color support

fmt_color.py

#!/usr/bin/env python3
import re
import sys
import string

MAX_LINE_LEN = 80
ANSI_ESCAPE_SEQUENCE = re.compile(r"\x1b\[[0-9;:]*?m")
ALL_ASCII_CHARS = {chr(i) for i in range(128)}
ALLOWED_CHARS = (ALL_ASCII_CHARS - set(string.whitespace)).union({" ", "\n"})

output = ""
printable_line_len = 0

for sequence_or_non_sequence in re.split(rf"({ANSI_ESCAPE_SEQUENCE.pattern})", sys.stdin.read()):
    for c in sequence_or_non_sequence:
        if c not in ALLOWED_CHARS:
            raise RuntimeError(f'character not allowed: "{c}"')
    if re.match(ANSI_ESCAPE_SEQUENCE, sequence_or_non_sequence):
        output += sequence_or_non_sequence
        continue
    for newline_or_word_or_whitespace in re.split(r"(\n|\b)", sequence_or_non_sequence):
        if newline_or_word_or_whitespace == "\n":
            output += "\n"
            printable_line_len = 0
        elif printable_line_len + len(newline_or_word_or_whitespace) > MAX_LINE_LEN:
            output += "\n" + newline_or_word_or_whitespace
            printable_line_len = len(newline_or_word_or_whitespace)
        else:
            output += newline_or_word_or_whitespace
            printable_line_len += len(newline_or_word_or_whitespace)

print(output)

I used nested re.split to iterate over each (ANSI escape sequence / newline / word / contiguous whitespace). Colors don't add to printable_line_len, newlines reset printable_line_len, and each word / whitespac is prepended with a newline whenever it crosses MAX_LINE_LEN. If a word / whitespace is prepended with a newline, printable_line_len becomes equal to the length of that word / whitespace. Of course this only works for characters that take up one space, so I made it so that only non-whitespace ascii characters, " " and "\n" are allowed.