Define a context manager generator
from contextlib import contextmanager
@contextmanager
def connect_to_db(address):
db = CrappyDBConnection(address)
try:
yield db
except ConnectionError:
logging.exception('Connection dropped')
db.cleanup('rollback')
else:
db.cleanup('commit')
db.disconnect()
# USE LIKE
with connect_to_db(address) as db:
...