wrappers.py

from tango import DevState, Except
from functools import wraps
import traceback

__all__ = ["only_in_states", "only_when_on", "fault_on_error"]

def only_in_states(allowed_states):
    """
      Wrapper to call and return the wrapped function if the device is
      in one of the given states. Otherwise a PyTango exception is thrown.
    """
    def wrapper(func):
        @wraps(func)
        def state_check_wrapper(self, *args, **kwargs):
            if self.get_state() in allowed_states:
                return func(self, *args, **kwargs)

            self.warn_stream("Illegal command: Function %s can only be called in states %s. Current state: %s" % (func.__name__, allowed_states, self.get_state()))
            Except.throw_exception("IllegalCommand", "Function can only be called in states %s. Current state: %s" % (allowed_states, self.get_state()), func.__name__)

        return state_check_wrapper

    return wrapper

def only_when_on():
    """
      Wrapper to call and return the wrapped function if the device is
      in the ON state. Otherwise None is returned and nothing
      will be called.
    """
    def inner(func):
        @wraps(func)
        def when_on_wrapper(self, *args, **kwargs):
            if self.get_state() == DevState.ON:
                return func(self, *args, **kwargs)

            return None

        return when_on_wrapper

    return inner

def fault_on_error():
    """
      Wrapper to catch exceptions. Sets the device in a FAULT state if any occurs.
    """
    def inner(func):
        @wraps(func)
        def error_wrapper(self, *args, **kwargs):
            try:
                return func(self, *args, **kwargs)
            except Exception as e:
                self.error_stream("Function failed.  Trace: %s", traceback.format_exc())
                self.Fault()
                return None

        return error_wrapper

    return inner