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, propagate_exc=False):
    """ ... """
    def wrap(fn):
        if db:
            try:
                return fn(db)
            except Exception, e:
                if propagate_exc:
                    raise
        else:
            tmpdb = env.get_db_cnx()
            try:
                res = fn(tmpdb)
                tmpdb.commit()
                return res
            except Exception, e:
                tmpdb.rollback()
                if propagate_exc:
                    raise
    return wrap

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

    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()
        return not self._propagate_exc


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

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

def do_rename_deco_prop(env, db=None):
    new_name = None

    @with_transaction(env, db, propagate_exc=True)
    def rename(db):
        db.execute()
        new_name = 'renamed' # local!
        return new_name
    new_name = rename
    
    if new_name:
        print '> it worked:', new_name
        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'

def do_rename_with_prop(env, db=None):
    new_name = None
    
    with transaction(env, db, propagate_exc=True) 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'
            do_rename(env)
            print
            print '-- testing explicit transaction'
            db = env.get_db_cnx()
            if do_rename(env, db) and do_rename(env, db):
                db.commit()
            else:
                db.rollback()

    for do_rename in (do_rename_deco_prop, do_rename_with_prop):
        env = Environment(True)
        print
        print '== (fail=%r, do_rename=%r)' % (fail, do_rename)
        print
        print '-- testing automatic transaction (propagate_exc)'
        try:
            do_rename(env)
        except ValueError:
            print 'oops'
        print
        print '-- testing explicit transaction (propagate_exc)'
        db = env.get_db_cnx()
        try:
            if do_rename(env, db) and do_rename(env, db):
                db.commit()
            else:
                db.rollback()
        except ValueError:
            print 'oops'
            db.rollback()