Search code examples
pythongeneratormysql-python

best way to run python generator cleanup code


I'm trying to write a generator function that gets rows out of a database and returns them one at a time. However, I'm not sure if the cleanup code marked ** below executes as I think it does. If it does not, what is the best way to put cleanup code inside a generator itself that executes after the last yield statement? I looked at catching StopIteration but that seems to be done from the caller, not within 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

Solution

  • You can use a context manager and the with statement. contextlib provides 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
    

    This will automatically call close() on dbc at the end of the with block, even if an exception has been raised.