清除尚未完全耗尽的迭代器

时间:2016-12-17 00:09:15

标签: python python-3.x generator python-3.5 paramiko

我对生成器的主要用法是处理存储在远程服务器上的CSV文件行。它允许我拥有一致的线性处理存储在其中的数据接口。

现在,我正在使用paramiko来访问存储文件的SFTP服务器 - 如果你没有关闭文件本身,paramiko有一个未正确关闭连接的突出问题。

我有一个访问sftp上单个文件的简单界面(这显然是伪代码 - 我省略了连接错误处理代码等等。)

 def sftp_read_file(filename):
       with paramiko.open(filename) as file_obj:
          for item in csv.reader(file_obj):
               yield item

 def csv_append_column(iter_obj, col_name, col_val):
     # header
     yield next(iter_obj) + (col_name, )
     for item in iter_obj:
         yield item + (col_val, )

假设我想通过运行有限数量的行的脚本来测试对文件执行的一组转换:

def main():
    for i, item in enumerate(csv_append_column(sftp_read_file('sftp://...'), 'A', 'B')):
        print(item)
        if i > 0 and i % 100 == 0:
            break

脚本将退出,但解释器永远不会在没有SIGINT的情况下终止。我可能的解决方案是什么?

1 个答案:

答案 0 :(得分:0)

这不是最优雅的解决方案,但也许我们可以通过将生成器包装在对象中来建立@ tadhg-mcdonald-jensen的建议:

class Stoppable(object):
    def __init__(self, fn):
        self.generator = fn

    def __enter__(self):
        return self.generator

    def __exit__(self, type_, value, traceback):
        self.generator.close()

然后像这样使用它:

def main():
    with Stoppable(sftp_read_file('sftp://...')) as reader:
        for i, item in enumerate(csv_append_column(reader, 'A', 'B')):
            print(item)
            if i > 0 and i % 100 == 0:
                break   

或者,如果我们不使用生成器方法进行流式传输,我们可以自行封装生成器:

def stopit(fn):
    rg = [ x for x in fn ]
    for x in rg:
        yield x

现在我们可以这样称呼:

def main():
    for i, item in enumerate(csv_append_column(stopit(sftp_read_file('...')), 'A', 'B')):
        print(item)
        if i > 0 and i % 100 == 0:
            break   

这将确保with block退出并paramiko关闭sftp连接,但代价是一次将所有行读入内存。