如何在Python中优雅地重定义input函数以按序返回预设值

在Python开发尤其是单元测试场景中，默认的input函数会阻塞程序等待用户手动输入，不利于自动化执行。重定义input函数让它按序返回预设值，可以很好地解决这个问题，下面介绍几种实现方案。

方案一：使用迭代器实现

迭代器可以依次返回元素，当元素耗尽时会抛出StopIteration异常，刚好符合按序返回预设值的需求。我们可以把预设值放到一个列表中，用iter函数生成迭代器，然后重定义input函数调用迭代器的__next__方法。

# 预设的输入值列表
preset_inputs = ["张三", "18", "男"]

# 生成迭代器
input_iter = iter(preset_inputs)

# 重定义input函数
def mock_input(prompt=""):
    # 原input函数会打印提示语，这里模拟相同行为
    if prompt:
        print(prompt, end="")
    return next(input_iter)

# 测试代码
print(mock_input("请输入姓名："))
print(mock_input("请输入年龄："))
print(mock_input("请输入性别："))

这种方式的优点是逻辑简单，预设值可以灵活修改，缺点是当预设值用尽后会直接抛出异常，需要提前确保预设值数量和调用次数匹配。

方案二：使用闭包封装状态

闭包可以保存外部函数的变量状态，我们可以用闭包来记录当前返回到第几个预设值，避免直接使用迭代器可能抛出的异常，还能自定义耗尽后的行为。

def create_mock_input(preset_values, default=None):
    index = 0
    def mock_input(prompt=""):
        nonlocal index
        if prompt:
            print(prompt, end="")
        if index < len(preset_values):
            value = preset_values[index]
            index += 1
            return value
        # 预设值耗尽后返回默认值，或者调用原input函数
        if default is not None:
            return default
        return input(prompt)
    return mock_input

# 预设值
preset_inputs = ["1001", "苹果", "5"]
# 创建模拟input函数
mock_input = create_mock_input(preset_inputs, default="0")

# 测试
print(mock_input("请输入商品编号："))
print(mock_input("请输入商品名称："))
print(mock_input("请输入购买数量："))
print(mock_input("请输入额外备注："))  # 预设值耗尽，返回默认值0

这种方式灵活性更高，可以自定义预设值耗尽后的处理逻辑，适合需要兼容部分真实输入的场景。

方案三：使用上下文管理器临时替换

如果只需要在某段代码执行期间重定义input函数，执行完后恢复原来的函数，使用上下文管理器是最优雅的方式，不会影响其他代码的正常input调用。

import builtins

class MockInputContext:
    def __init__(self, preset_values):
        self.preset_values = preset_values
        self.original_input = None
        self.index = 0

    def __enter__(self):
        # 保存原input函数
        self.original_input = builtins.input
        # 定义模拟函数
        def mock_input(prompt=""):
            if prompt:
                print(prompt, end="")
            if self.index < len(self.preset_values):
                value = self.preset_values[self.index]
                self.index += 1
                return value
            # 预设值耗尽调用原函数
            return self.original_input(prompt)
        # 替换内置input
        builtins.input = mock_input
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 恢复原input函数
        builtins.input = self.original_input

# 测试代码
preset_inputs = ["2024", "10", "1"]
with MockInputContext(preset_inputs):
    print(input("请输入年份："))
    print(input("请输入月份："))
    print(input("请输入日期："))

# 上下文外恢复原函数，正常等待输入
# print(input("请输入测试内容："))

这种方式的好处是替换是临时的，不会对全局环境造成长期影响，非常适合单元测试场景，测试完成后自动恢复原input函数，避免影响其他测试用例。

注意事项

• 重定义input函数时尽量模拟原函数的行为，比如原input函数会打印提示语且不换行，自定义函数也要保持相同行为，避免影响原有代码逻辑。
• 如果是全局替换input函数，要注意在不需要的时候恢复原函数的引用，避免影响其他模块的正常使用。
• 预设值的数量要和input调用次数匹配，或者做好耗尽后的处理逻辑，避免出现意外的异常。
• 在单元测试中优先使用上下文管理器的方式，隔离性更好，不会造成测试用例之间的相互影响。

方案对比

Pythoninput函数函数重定义预设值返回单元测试修改时间：2026-07-05 13:42:25