def switch_modal_function(stub=None, use_modal=True, **kwargs): def wrapper(inner_func): if use_modal: assert stub is not None, "If using modal, please provide `stub`." return stub.function(**kwargs)(inner_func) else: return inner_func return wrapper