How to return a generator in Python

Question

I am thinking of designing my function to return a generator in combination with a database query. but having some question on the concept of iterators

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

The question is, when i iterate the generator in the main function, i hit "error on closed cursor". Should i close or not close it in the func()? I suppose when the call to func() end, the dbconn variable will be out of scope and i shouldn't have to worry about closing?

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

How should i design this? returing a data structure like a list would be better? thanks

Answer 1

You could use a Context Manager , eg (contains some pseudocode):

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

Of course this is only useful if execute(sql) really returns a generator. If you put all data into a list (and thus into memory) before closing the connection, your problem will be obsolete.

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

---

In response to your comment:

If your database adapter follows the python DB API spec, an efficient way is to fetch a bunch of rows with fetchmany multiple times.

The following code fetches rows in chunks of 100, and explicitly calls dbclose when the execution leaves the with block:

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

I don't have much experience on using database, but I think you should retrieve the query results and return it as a list. If you really need an iterator (but I cant see why), then return an iterator over the list 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 

Now, if their exist a way to retrieve the nth element of a query, something like execute(sql, n) which return None if n is too big, then you could use a 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()

Now, this is not what I would recommend, mainly because the connection to the db stays open while the iterator is not finished. And it might never happens if something fails or is badly designed.

Answer 3

You cannot try to manipulate a cursor after you have closed the database connection I will try with this approach:

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/

I hope it helps