from __future__ import with_statement # this requires Python 2.5 or 2.6

class Connection(object):
    def __init__(self, fail):
        self._fail = fail
    def execute(self):
        if self._fail:
            print '>> arghh!'
            a = int('')
        else:
            print '>> execute something successfully...'
    def commit(self):
        print '>> commit transaction'
    def rollback(self):
        print '>> rollback transaction'

class Environment(object):
    def __init__(self, fail):
        self._fail = fail
    def get_db_cnx(self):
        print ">> Environment.get_db_cnx"
        return Connection(self._fail)


def with_transaction(env, db):
    """ ... """
    def wrap(fn):
        if db:
            fn(db)
        else:
            tmpdb = env.get_db_cnx()
            try:
                fn(tmpdb)
                tmpdb.commit()
            except Exception:
                tmpdb.rollback()
                raise
    return wrap


class transaction(object):
    def __init__(self, env, db):
        self._env = env
        self._db = db

    def __enter__(self):
        if not self._db:
            self._db = self._env.get_db_cnx()
            self._env = None
        return self._db

    def __exit__(self, et, ev, tb):
        if self._env is None:
            if et is None:
                self._db.commit()
            else:
                self._db.rollback()


# Trac 0.12
def do_rename_deco(env, db=None):
    new_name = [None]

    @with_transaction(env, db)
    def rename(db):
        db.execute()
        new_name[0] = 'renamed'
    
    if new_name[0]:
        print '> it worked:', new_name[0]
        return True
    else:
        print '> rename failed'


# Trac 0.13
def do_rename_with(env, db=None):
    new_name = None
    
    with transaction(env, db) as db:
        db.execute()
        new_name = 'renamed' # modifies `new_name` above
    
    if new_name:
        print '> it worked:', new_name
        return True
    else:
        print '> rename failed'



if __name__ == '__main__':
    for do_rename in (do_rename_deco, do_rename_with):
        for fail in (False, True):
            env = Environment(fail)
            print
            print '== (fail=%r, do_rename=%r)' % (fail, do_rename)
            print
            print '-- testing automatic transaction'
            try:
                do_rename(env)
            except:
                pass
            print
            print '-- testing explicit transaction'
            db = env.get_db_cnx()
            try:
                if do_rename(env, db) and do_rename(env, db):
                    db.commit()
                else:
                    db.rollback()
            except:
                db.rollback()