Python: Write colored text in file

Question

I want to write a file containing some arithmetic problems for my little child. I want to have some coloring, so that he can easily make the difference between plus and minus. This worked for me very well. Unfortunately, only in a terminal.

import random as rd
from termcolor import colored

N = 10
MAX = 100
f = open("math.txt", "w")


def get_random_str():

    a = rd.randint(1, MAX)
    b = rd.randint(1, MAX)

    if a < MAX*0.4:
        string = "%3d "%a + str(colored('+', 'blue')) + " %d = \n"%(b)

    else:
        if a>b:
            string = "%3d "%a + str(colored('-', 'red')) + " %d = \n"%(b)

        else:
            string = "%3d "%a + str(colored('-', 'red')) + " %d = \n"%(b)

    return string
#-------------------------------------------------------------------------

for i in range(1,N):
    print i, get_random_str()

When I try to write the output in a file, of course I just get the color codes e.g. "[34m+[0m" instead of a red "-" and a blue "+".

Any idea how to solve this task?

Answer 1

You could check out Pygments with any suitable lexer and a TerminalFormatter.

E.g. the following code:

import sys
from pygments import highlight
from pygments.formatters.terminal import TerminalFormatter
from pygments.lexer import RegexLexer
from pygments.token import Token


class ArithmeticLexer(RegexLexer):
    tokens = {
        'root': [
            (r'[ \n]', Token.Whitespace),
            (r'\d+', Token.Number),
            (r'\+', Token.Plus),
            (r'-', Token.Minus),
            (r'\*', Token.Multiply),
            (r'/', Token.Division),
        ]
    }

COLOR_SCHEME = {
    Token.Whitespace: ('', ''),
    Token.Number: ('darkgreen', 'green'),
    Token.Plus: ('darkred', 'red'),
    Token.Minus: ('darkblue', 'blue'),
    Token.Multiply: ('darkyellow', 'yellow'),
    Token.Division: ('brown', 'fushia'),
}

if __name__ == '__main__':
    with open(sys.argv[1], 'rb') as f:
        for line in f:
            line = highlight(line, ArithmeticLexer(), TerminalFormatter(colorscheme=COLOR_SCHEME))
            print line.strip()

Gives:

When ran using file with given contents. The usage is <script_name> <input_file_name>.

The colors' reference. The colors in COLOR_SCHEME are tuples of (lightscheme, darkscheme). By defaults TerminalFormatter uses lightscheme.