diff --git a/RCUSCC/RCUSCC/wrappers.py b/RCUSCC/RCUSCC/wrappers.py index dfd4fe3f444c0cbb572e53f7b0386b2882171e55..cb557c4183abcaa812b24b4dbe721290931775e4 100644 --- a/RCUSCC/RCUSCC/wrappers.py +++ b/RCUSCC/RCUSCC/wrappers.py @@ -6,33 +6,31 @@ __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. + 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() not in allowed_states: - return None + if self.get_state() in allowed_states: + return func(self, *args, **kwargs) - return func(self, *args, **kwargs) + return None 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. + 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 None + if self.get_state() == DevState.ON: + return func(self, *args, **kwargs) - return func(self, *args, **kwargs) + return None return when_on_wrapper @@ -40,7 +38,6 @@ 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: