如何在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塊時顯式調用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)

Answer 2

我在使用數據庫方面經驗不足，但是我認為您應該檢索查詢結果並將其作為列表返回。 如果您確實需要一個迭代器（但我看不到為什么），則在列表ret返回一個迭代器：

def func():
    sql =" select some rows "
    dbconn = "connect and open to dtabase code"
    ret = execute(sql)              # a list
    dbclose()
    return (elmt for elmt in ret)   # return an iterator over ret 

現在，如果它們存在一種檢索查詢的第n個元素的方法，例如execute(sql, n) ，如果n太大則返回None ，則可以使用yield：

 def func():
    sql =" select some rows "
    dbconn = "connect and open to dtabase code"

    n = 0
    ret = execute(sql,n)    # return the n-th element
    while ret is not None:
        yield ret
        n += 1
        ret = execute(sql,n)

    dbclose()

現在，這不是我所建議的，主要是因為在迭代器未完成時與數據庫的連接保持打開狀態。 如果某件事失敗或設計不當，它可能永遠不會發生。

Answer 3

關閉數據庫連接后，您將無法嘗試操作游標，我將嘗試使用這種方法：

def func(params):
    sql = "query to execute"
    cursor = execute(sql, params)
    return cursor.fetchall() # retrieves all posible results as a sequence of sequences,
                             # i.g. list of tuples(*)

### Main ###
# Open database connection
# Create cursor
for elem in func(): # Call to retrieve desired element's method and do something with that
    # Do something
# Close cursor
# Close database connection

（*） http://www.python.org/dev/peps/pep-0249/

希望對您有所幫助