Best way to run python generator cleanup code

I am trying to write a generator function that outputs rows from a database and returns them one at a time. However, I'm not sure if the cleanup code marked ** below is executing as I think. If this is not the case, what is the best way to put the cleanup code inside the generator itself, which runs after the last yield statement? I looked at the StopIteration trap, but it seems to be done from the caller, not inside the generator.

def MYSQLSelectGenerator(stmt):
...
try:   
    myDB = MySQLdb.connect(host=..., port=..., user=..., passwd=..., db=...)   
    dbc=myDB.cursor()
    dbc.execute(stmt)
    d = "asdf"
    while d is not None:
        d = dbc.fetchone() #can also use fetchmany() to be more efficient
        yield d
    dbc.close() #** DOES THIS WORK AS I INTEND, MEANING AS SOON AS d = "None"
except MySQLdb.Error, msg:
    print("MYSQL ERROR!")
    print msg
+4
source share
3 answers

You can use the context manager and operator with. provides : contextlib closing

from contextlib import closing

myDB = MySQLdb.connect(host=..., port=..., user=..., passwd=..., db=...)   
with closing(myDB.cursor()) as dbc:
    dbc.execute(stmt)
    d = "asdf"
    while d is not None:
        d = dbc.fetchone() #can also use fetchmany() to be more efficient
        yield d

close() dbc with, .

+1

dbc.close(), d is None, , . finally . , , dbc.close(), :

try:   
    myDB = MySQLdb.connect(host=..., port=..., user=..., passwd=..., db=...)   
    dbc = myDB.cursor()
    dbc.execute(stmt)
    d = "asdf"
    while d is not None:
        d = dbc.fetchone() #can also use fetchmany() to be more efficient
        yield d
except MySQLdb.Error, msg:
    print("MYSQL ERROR!")
    print msg
finally:
    dbc.close()
+3

One thing you can do is use a sentence finally. Another option (which may be redundant, but useful for familiarization) is to create a class that works with the operator with:

class DatabaseConnection:
    def __init__(self, statement):
        self.statemet = statement
    def __enter__(self): 
        self.myDB = MySQLdb.connect(host=..., port=...,user=...,passwd=...,db=...)
        self.dbc = myDB.cursor()
        self.dbc.execute(self.statement)
        self.d = "asdf"
    def __exit__(self, exc_type, exc_value, traceback):
        self.dbc.close()

    def __iter__(self):
        while self.d is not None:
            self.d = self.dbc.fetchone()
            yield self.d


with DatabaseConnection(stmnt) as dbconnection:
    for i in dbconnection:
        print(i)
+3
source

All Articles