Python `doctest` with ANSI escapes

Question

I have some code for which I have some doctest, e.g.

def foo(x):
    """
    Examples:
        >>> s = foo()
        >>> print(repr(s))
        '\x1b[0;35mfoo\x1b[0m \x1b[0;36mbar\x1b[0m'
        >>> print(s)
        foo bar
    """
    return '\u001b[0;35mfoo\u001b[0m \u001b[0;36mbar\u001b[0m'

Of course, in a ANSI-escape capable terminal, the last print() is actually colored.

However, if I run doctest here, I will get a failing test. How could I solve this?

I was hoping either to have a way of inserting ANSI escapes or, preferably, a doctest directive for ignoring them, but I could not find any information on this in the docs.

norok2 · Accepted Answer · 2019-08-08T10:04:11.617

A silly workaround to this is to use a in-code comparison instead of relying on doctests comparisons, e.g.:

def foo(x):
    """
    Examples:
        >>> s = foo()
        >>> print(repr(s))
        '\x1b[0;35mfoo\x1b[0m \x1b[0;36mbar\x1b[0m'
        >>> print(s == '\x1b[0;35mfoo\x1b[0m \x1b[0;36mbar\x1b[0m')
        True
    """
    return '\u001b[0;35mfoo\u001b[0m \u001b[0;36mbar\u001b[0m'

Another possibility would be to filter out ANSI escapes either right inside print():

import re


def remove_ansi_escapes(text):
    ansi_escape = re.compile(r'(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]')
    return ansi_escape.sub('', text)


def foo(x):
    """
    Examples:
        >>> s = foo()
        >>> print(repr(s))
        '\x1b[0;35mfoo\x1b[0m \x1b[0;36mbar\x1b[0m'
        >>> print(remove_ansi_escapes(s))
        foo bar
    """
    return '\u001b[0;35mfoo\u001b[0m \u001b[0;36mbar\u001b[0m'

or by capturing STDOUT if the printing is happening deeper in the code:

from contextlib import redirect_stdout
import io


def foo(x):
    """
    Examples:
        >>> buffer = io.StringIO()
        >>> with redirect_stdout(buffer):
        ...     foo()
        >>> print(remove_ansi_escapes(buffer.getvalue()))
        foo bar
    """
    s = '\u001b[0;35mfoo\u001b[0m \u001b[0;36mbar\u001b[0m'
    print(s)

Python `doctest` with ANSI escapes

1 Answers1