wrappers.py

from tango import DevState
from functools import wraps
import traceback

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

def only_in_states(func, allowed_states):
    """
      Wrapper to call and return the wrapped function if the device is
      in one of the given states. Otherwise None is returned and nothing
      will be called.
    """
    @wraps(func)
    def state_check_wrapper(self, *args, **kwargs):
        if self.get_state() in allowed_states:
            return func(self, *args, **kwargs)

        return None

    return state_check_wrapper

def only_when_on(func):
    """
      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.
    """
    @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

def fault_on_error(func):
    """
      Wrapper to catch exceptions. Sets the device in a FAULT state if any occurs.
    """
    @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