如何为输出彩色文本的函数编写Python doctest测试？

你遇到的问题本质是：print_result_in_color()的实际输出包含ANSI颜色转义序列（比如Fore.GREEN对应的\x1b[32m），终端显示时会解析这些序列变成颜色，但doctest对比的是原始字符串内容，所以看起来输出都是"Hello world"，实际底层字符串不一样，导致测试失败。

下面提供两种可行的解决方案：

方案1：忽略颜色，只验证文本内容 #

通过捕获函数的输出，用正则去掉所有ANSI转义序列，再和预期文本对比：

import sys
import re
from colorama import Fore
from io import StringIO

# 定义去除ANSI转义序列的工具函数
def strip_ansi(text):
    ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[.*?[a-zA-Z])')
    return ansi_escape.sub('', text)

def print_result():
    """Prints out the result.

    >>> print_result()
    Hello world
    """
    print('Hello world')

def print_result_in_color():
    """Prints out the result in green.

    >>> captured_output = StringIO()
    >>> sys.stdout = captured_output
    >>> print_result_in_color()
    >>> sys.stdout = sys.__stdout__  # 恢复标准输出
    >>> strip_ansi(captured_output.getvalue()).strip()
    'Hello world'
    """
    print(Fore.GREEN + 'Hello world' + Fore.RESET)

if __name__ == '__main__':
    import doctest
    doctest.testmod()

方案2：在docstring中包含颜色转义信息 #

直接在预期结果里写入Fore.GREEN和Fore.RESET对应的原始转义字符串，让doctest精准匹配完整输出：

from colorama import Fore

def print_result():
    """Prints out the result.

    >>> print_result()
    Hello world
    """
    print('Hello world')

def print_result_in_color():
    """Prints out the result in green.

    >>> print_result_in_color()
    \x1b[32mHello world\x1b[0m
    """
    print(Fore.GREEN + 'Hello world' + Fore.RESET)

if __name__ == '__main__':
    import doctest
    doctest.testmod()

注：\x1b[32m是Fore.GREEN对应的ANSI转义码，\x1b[0m是Fore.RESET的转义码，直接写入后doctest就能正确匹配带颜色的输出字符串。

内容的提问来源于stack exchange，提问作者freude