带有 ANSI 转义的 Python `doctest`答案

【问题标题】：Python `doctest` with ANSI escapes带有 ANSI 转义的 Python `doctest`
【发布时间】：2019-11-19 07:08:46
【问题描述】：

我有一些代码，我有一些 doctest，例如

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'

当然，在支持 ANSI 转义的终端中，最后一个 print() 实际上是彩色的。

但是，如果我在这里运行doctest，我会得到一个失败的测试。我该如何解决这个问题？

我希望有一种插入 ANSI 转义符的方法，或者最好是一个忽略它们的 doctest 指令，但我在 docs 中找不到任何相关信息。

【问题讨论】：

标签： python doctest ansi-escape

【解决方案1】：

一个愚蠢的解决方法是使用代码内比较而不是依赖于 doctests 比较，例如：

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'

另一种可能性是filter out ANSI escapes 就在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'

或capturing STDOUT，如果打印发生在代码的更深处：

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)

【讨论】：