Transaction based Python script for PSSE

transactions.py

"""
Transaction
-----------

Wrap your PSSe  Python code in this transaction and watch it rollback
when the case blows up. 
"""
from __future__ import with_statement
import contextlib
import os
import sys
import tempfile

#----------------------------------------------------------------------------------
# Some junk to import PSSE put in a library!

PSSE_VERSION = 32
PSSE_BASE = os.path.join('c:\\', 'program files', 'pti', 'psse%d' % PSSE_VERSION)
PSSE_LOCATION = os.path.join(PSSE_BASE, 'pssbin')
PSSE_EXAMPLES = os.path.join(PSSE_BASE, 'example')

sys.path.append(PSSE_LOCATION)
os.environ['PATH'] = os.environ['PATH'] + ';' + PSSE_LOCATION

import psspy
psspy.throwPsseExceptions = True


#----------------------------------------------------------------------------
# Transactions.
class TransactionError(psspy.PsseException):
    pass

@contextlib.contextmanager
def transaction(validator=None):
    """
    Begin by saving the current state to a saved
    file. Then read back from that saved file if the validator test
    at the end fails.

    Remember you must solve your case during the transaction. Otherwise
    many of your changes will not have a chance to affect the system yet.

    Arguments:
        validator - A function without arguments to call that will return True
            if the case is valid. Otherwise the case will be rolled back.

            if no validator is supplied, the case will always be rolled back.

    >>> with transaction():
            add_generator()
            # remember to solve your case.
            psspy.fnsl()

    """

    # create a temporary file securely on the operating system.
    (file_handle, temp_save_filename) = tempfile.mkstemp(suffix=".sav")
    os.close(file_handle)
    try:
        ierr = psspy.save(temp_save_filename)
        if ierr: raise TransactionError('Couldnt psspy.save %d' % ierr)

        yield

        # if validator == True don't rollback.
        if validator and validator():
            return

        # ok lets rollback because validator is None or False.
        ierr = psspy.case(temp_save_filename)
        if ierr: raise TransactionError('Couldnt psspy.case %d' % ierr)

    finally:

        try:
            os.remove(temp_save_filename)
        except OSError:
            pass