python - 如何在 Python 中返回生成器

Question

我正在考虑设计我的函数来返回与数据库查询相结合的生成器。但对迭代器的概念有一些疑问

def func():
    sql =" select some rows "
    dbconn = "connect and open to dtabase code"
    ret = ( execute(sql)  ) <----- returning a generator?
    dbclose <----  I close the db connection here, but it gives me error
    return ret

问题是，当我在主函数中迭代生成器时，我遇到了“关闭光标时出错”。我应该在 func() 中关闭还是不关闭它？我想当对 func() 的调用结束时， dbconn 变量将超出范围，我不必担心关闭？

 # main function
 for it in func():
     do something with it
 close dbconn here?

我应该如何设计这个？返回像列表这样的数据结构会更好吗？谢谢

Answer 1

您可以使用Context Manager ，例如(包含一些伪代码):

from contextlib import contextmanager

@contextmanager
def func():
    sql =" select some rows "
    dbconn = "connect and open to dtabase code"
    yield execute(sql)  # returns this really a generator?
    dbclose #pseudocode, you probably want to put this in a try/finally block

with func() as result:
    for it in result:
         do something with it

当然，这仅在 execute(sql) 确实返回生成器时才有用。如果您在关闭连接之前将所有数据放入列表(从而放入内存)，那么您的问题将被淘汰。

def func():
    sql =" select some rows "
    dbconn = "connect and open to dtabase code"
    ret = list( execute(sql)  ) 
    dbclose # no problem here, since all data is already fetched
    return ret

<小时/>

回复您的评论:

如果您的数据库适配器遵循 python DB API 规范，则一种有效的方法是使用 fetchmany 多次获取一堆行。

以下代码以 100 为一组获取行，并在执行离开 with block 时显式调用 dbclose:

def cursor_iter(cursor, num_of_rows=100):
    while True:
        rows = cursor.fetchmany(num_of_rows)
        if not rows: break
        for row in rows:
            yield row

@contextmanager
def func():
    sql = "select some rows"
    dbconn = connect_and_open_database()
    cursor = dbconn.cursor()
    cursor.execute(sql)
    yield cursor_iter(cursor)
    dbclose()

with func() as result:
    for row in result: 
        do_something(row)