如何测试含input和print的字符串形式Python源代码？

问题背景 #

我正在开发一个简易编程测验应用，要求用户提交Python源代码作答。比如一道求和题：

题目要求：给定n个数值，求和并输出
输入格式：

N
a[0], a[1], ..., a[N-1]

输出格式：

s

示例输入：

3
1 3 5

示例输出：

9

用户提交的代码是字符串形式，比如：

N = int(input())
print(sum([int(x) for x in input().split()]))

服务器接收字符串形式的源代码后，需要自动测试代码的正确性。目前我尝试用exec执行代码，但不知道怎么程序化给input()提供测试数据，不过已经了解可以重定向stdout捕获print的输出，求完整解决方案。

解决方案 #

1. 核心思路 #

通过替换sys.stdin和sys.stdout实现输入模拟与输出捕获：

• 用io.StringIO构造模拟输入流，将测试数据写入其中，替代系统标准输入
• 用io.StringIO接管系统标准输出，捕获print语句的输出内容
• 执行用户代码前后保存并恢复原始标准输入输出，避免影响后续逻辑

2. 完整实现代码 #

import sys
from io import StringIO

def test_user_code(source_code, test_input, expected_output):
    # 保存原始的标准输入输出
    original_stdin = sys.stdin
    original_stdout = sys.stdout
    
    try:
        # 构造模拟输入流
        sys.stdin = StringIO(test_input)
        # 构造输出捕获流
        sys.stdout = captured_output = StringIO()
        
        # 执行用户提交的代码
        exec(source_code, globals())
        
        # 获取捕获的输出，去除首尾空白（处理多余换行/空格）
        actual_output = captured_output.getvalue().strip()
        
        # 验证结果是否匹配
        assert actual_output == expected_output, f"预期输出: {expected_output}, 实际输出: {actual_output}"
        return True
    except SyntaxError as e:
        print(f"代码语法错误: {str(e)}")
        return False
    except Exception as e:
        print(f"代码运行错误: {str(e)}")
        return False
    finally:
        # 恢复原始的标准输入输出
        sys.stdin = original_stdin
        sys.stdout = original_stdout

# 测试示例
if __name__ == "__main__":
    user_code = """
N = int(input())
print(sum([int(x) for x in input().split()]))
"""
    test_input = "3\n1 3 5"
    expected_output = "9"
    
    test_result = test_user_code(user_code, test_input, expected_output)
    print("测试通过" if test_result else "测试失败")

3. 关键细节说明 #

• 输入模拟：将测试数据按换行符拼接成字符串，input()会依次读取每一行内容，完全模拟用户手动输入的逻辑
• 输出处理：通过strip()去除输出首尾的空白字符，避免因用户代码多打换行/空格导致的匹配失败
• 异常分类：区分语法错误和运行时错误，方便给用户返回更精准的错误提示
• 环境隔离：在finally块强制恢复原始输入输出，防止后续代码的标准流被篡改

4. 额外注意事项 #

• 需限制用户代码的执行时间和资源占用，可借助timeout-decorator等工具实现超时控制，避免恶意代码拖垮服务器
• 可准备多组测试用例，循环调用测试函数，覆盖更多边界场景（比如N=1、负数求和等）
• 若题目允许输出格式有微小差异（比如末尾多一个换行），可调整输出对比逻辑，比如只比较内容是否一致

内容的提问来源于stack exchange，提问作者Wt.N