Define a context manager class

class_contextmanager.py

class DBConnection(object):
    def __init__(self, address):
        self.crappy_db = CrappyDBConnection(address)
    
    def __enter__(self):
        self.crappy_db.connect()
        return self  # what is returned here will be availabe via "as"
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type == ConnectionError:
            logging.exception('Connection dropped')
            self.crappy_db.cleanup('rollback')
        else:
            self.crappy_db.cleanup('commit')
        self.creappy_db.disconnect()
    
    def more_awesome_methods(self):
        ...


# THIS
with DBConnection('127.0.0.1') as db:
    ...


# INSTEAD OF
db = CrappyDBConnection('127.0.0.1')
try:
    db.connect()
    ...
except ConnectionError:
    logging.exception('Connection dropped')
    db.cleanup('rollback')
else:
    db.cleanup('commit')
finally:
    db.disconnect()