如何使StringIO作为程序stdin时，input等待线程写入输入内容

问题描述 #

我想把程序的stdin重定向到StringIO()对象，用来模拟用户对input()语句的响应。目前的代码是这样的：

from io import StringIO
import sys

newstdin = StringIO()
sys.stdin = newstdin
newstdin.write("hey")
newstdin.seek(0)
response = input()
print(response)

当StringIO()里有内容的时候，代码能正常运行；但如果对象是空的，程序直接抛出EOF错误，而不是像原生sys.stdin那样等待输入。我想实现的是让input()语句能等待另一个线程向StringIO()对象写入内容，请问该怎么做？

解决方案 #

这个问题我之前也碰到过——StringIO本质就是个纯内存的缓冲区，它没有内置的阻塞等待机制，一旦读到末尾就直接返回EOF，这就是空缓冲区时报错的原因。要实现阻塞等待其他线程写入的效果，我们需要自己封装一个类似文件对象的类，用线程同步的原语来处理等待逻辑。

下面是一个可行的实现：

from io import StringIO
import sys
import threading
from time import sleep

class BlockingStringIO:
    def __init__(self):
        self.buffer = StringIO()
        self.lock = threading.Lock()
        self.condition = threading.Condition(self.lock)
        self.closed = False

    def write(self, data):
        with self.lock:
            self.buffer.write(data)
            self.condition.notify_all()  # 通知所有等待的线程：有新数据了

    def readline(self):
        with self.lock:
            # 循环等待，直到缓冲区有数据或者对象被关闭
            while True:
                self.buffer.seek(0, 2)  # 移动到缓冲区末尾，检查是否有新数据
                if self.buffer.tell() > 0 or self.closed:
                    break
                self.condition.wait()  # 没有数据就进入等待状态
                
            self.buffer.seek(0)
            line = self.buffer.readline()
            # 把读取后剩下的内容重新存入缓冲区，方便下一次读取
            remaining_content = self.buffer.read()
            self.buffer = StringIO(remaining_content)
            return line if line else None

    def close(self):
        with self.lock:
            self.closed = True
            self.condition.notify_all()  # 通知等待的线程：对象已关闭

# 测试代码
def delayed_writer(io_obj):
    sleep(2)  # 模拟2秒的延迟写入
    io_obj.write("Hello from background thread!\n")

if __name__ == "__main__":
    blocking_io = BlockingStringIO()
    sys.stdin = blocking_io

    # 启动后台写入线程
    threading.Thread(target=delayed_writer, args=(blocking_io,), daemon=True).start()

    print("Waiting for input...")
    response = input()  # 这里会阻塞，直到后台线程写入数据
    print(f"Got response: {response.strip()}")

代码说明 #

• 我们封装的BlockingStringIO类内部用StringIO存储数据，同时用Condition条件变量实现阻塞等待逻辑——这是Python线程同步里用来处理“等待某个条件成立”场景的常用工具。
• write方法写入数据后，会立刻通知所有处于等待状态的线程，告诉它们可以尝试读取数据了。
• readline方法（input()底层会调用这个方法）会先检查缓冲区是否有数据：如果没有，就进入等待状态；一旦收到新数据的通知，就读取一行内容，并把剩余的内容保留在缓冲区中供后续读取。
• 如果调用close()方法，会标记对象为关闭状态，并通知所有等待的线程结束等待。

这样修改后，当input()被调用时，如果缓冲区是空的，它会像原生stdin一样阻塞等待，直到另一个线程写入数据为止。

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