diff --git a/RCUSCC/RCUSCC/wrappers.py b/RCUSCC/RCUSCC/wrappers.py index cb557c4183abcaa812b24b4dbe721290931775e4..14d8f1d6426bee6fd1a0f5b164d30ebd8d07d698 100644 --- a/RCUSCC/RCUSCC/wrappers.py +++ b/RCUSCC/RCUSCC/wrappers.py @@ -4,20 +4,23 @@ import traceback __all__ = ["only_in_states", "only_when_on", "fault_on_error"] -def only_in_states(func, allowed_states): +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 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) + def wrapper(func): + @wraps(func) + def state_check_wrapper(self, *args, **kwargs): + if self.get_state() in allowed_states: + return func(self, *args, **kwargs) - return None + return None + + return state_check_wrapper - return state_check_wrapper + return wrapper def only_when_on(func): """