如何明确要求调用方关闭生成器，避免未完全消费导致SSH连接泄漏？

我在编程时用生成器函数管理SSH连接资源，目的是避免两次遍历文件名列表，示例代码如下：

# 用生成器避免两次遍历文件名列表
def _yield_fname(host_address, usr, pwd, datapath):
    with paramiko.SSHClient() as ssh_client:
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh_client.connect(host, username=usr, password=pwd)
        
        stdin, stdout, stderr = ssh_client.exec_command(
            'find . -type d -maxdepth 1 -printf "%f\n"'
        )
        for line in stdout:
            yield fname

def get_foos():
    foos= []
    for fname in self._yield_fname(self.FOO_PATH):
        foos.append(Foo.from_fname(fname))
    return foos

def get_bars():
    bars= []
    for fname in self._yield_fname(self.FOO_PATH):
        bars.append(Bar.from_fname(fname))
    return bars

目前我会完全消费生成器来确保连接关闭，但担心其他人（比如未来的自己）没有完全消费生成器，导致SSH连接无法正常关闭。请问如何明确提示调用方，即使未完全消费生成器，也必须在使用后将其关闭？

可行的解决方案 #

1. 将生成器包装为上下文管理器 #

让生成器支持with语句，强制调用方通过上下文管理资源，无论是否完全消费，退出上下文时都会自动清理SSH连接。可以用contextlib.contextmanager装饰器改造：

from contextlib import contextmanager

@contextmanager
def _yield_fname(host_address, usr, pwd, datapath):
    with paramiko.SSHClient() as ssh_client:
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh_client.connect(host, username=usr, password=pwd)
        
        stdin, stdout, stderr = ssh_client.exec_command(
            'find . -type d -maxdepth 1 -printf "%f\n"'
        )
        # 转换为迭代器，方便后续处理剩余输出
        fname_iter = (line.strip() for line in stdout)
        try:
            yield fname_iter
        finally:
            # 读取剩余输出，避免SSH连接因未处理完数据流而阻塞
            for _ in fname_iter:
                pass

调用方式变为：

def get_foos():
    foos = []
    with self._yield_fname(self.FOO_PATH) as fname_iter:
        for fname in fname_iter:
            foos.append(Foo.from_fname(fname))
    return foos

2. 利用生成器关闭机制+文档提示 #

在生成器内部处理GeneratorExit异常，确保生成器被手动关闭时也能清理资源，同时在函数文档中明确提示调用方规则：

def _yield_fname(host_address, usr, pwd, datapath):
    """
    生成远程主机指定路径下的目录名。
    注意：若未完全迭代生成器，必须手动调用`close()`方法，否则SSH连接可能无法正常关闭。
    """
    ssh_client = paramiko.SSHClient()
    try:
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh_client.connect(host, username=usr, password=pwd)
        
        stdin, stdout, stderr = ssh_client.exec_command(
            'find . -type d -maxdepth 1 -printf "%f\n"'
        )
        for line in stdout:
            yield line.strip()
    except GeneratorExit:
        # 生成器被关闭时，读取剩余输出释放连接
        for _ in stdout:
            pass
    finally:
        ssh_client.close()

3. 封装迭代逻辑，隐藏生成器细节 #

不直接暴露生成器给调用方，而是提供封装好的处理接口，内部统一确保资源被正确清理：

def _yield_fname(host_address, usr, pwd, datapath):
    # 原生成器实现
    with paramiko.SSHClient() as ssh_client:
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh_client.connect(host, username=usr, password=pwd)
        
        stdin, stdout, stderr = ssh_client.exec_command(
            'find . -type d -maxdepth 1 -printf "%f\n"'
        )
        for line in stdout:
            yield line.strip()

def process_fnames(self, handler):
    """
    批量处理远程路径下的目录名，handler为单个目录名的处理函数
    """
    for fname in self._yield_fname(self.FOO_PATH):
        handler(fname)

def get_foos(self):
    foos = []
    self.process_fnames(lambda fname: foos.append(Foo.from_fname(fname)))
    return foos

def get_bars(self):
    bars = []
    self.process_fnames(lambda fname: bars.append(Bar.from_fname(fname)))
    return bars

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