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 return None when the device isn't in ON state.

      If in ON state, calls & returns the wrapped function.
    """

    @wraps(func)
    def state_check_wrapper(self, *args, **kwargs):
        if self.get_state() not in allowed_states:
            return None

        return func(self, *args, **kwargs)

    return state_check_wrapper

def only_when_on(func):
    """
      Wrapper to return None when the device isn't in ON state.

      If in ON state, calls & returns the wrapped function.
    """

    @wraps(func)
    def when_on_wrapper(self, *args, **kwargs):
        if self.get_state() != DevState.ON:
            return None

        return func(self, *args, **kwargs)

    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