from functools import wraps from ...messaging import RollbotResponse, RollbotFailure def require_min_args(n, alert_response=None): def decorator(cls): old_on_command = cls.on_command @wraps(old_on_command) def wrapper(self, db, message): if message.raw_args is None or len(message.arg_list()) < n: failure = RollbotFailure.INVALID_ARGUMENTS.with_reason(alert_response or f"{self.command.title()} requires at least {n} argument(s)") return RollbotResponse(message, failure=failure, debugging=failure.get_debugging()) return old_on_command(self, db, message) setattr(cls, "on_command", wrapper) return cls return decorator # TODO refactor to share code? def require_args(n, alert_response=None): def decorator(cls): old_on_command = cls.on_command @wraps(old_on_command) def wrapper(self, db, message): if message.raw_args is None or len(message.arg_list()) != n: failure = RollbotFailure.INVALID_ARGUMENTS.with_reason(alert_response or f"{self.command.title()} requires exactly {n} argument(s)") return RollbotResponse(message, failure=failure, debugging=failure.get_debugging()) return old_on_command(self, db, message) setattr(cls, "on_command", wrapper) return cls return decorator def require_admin(cls): old_on_command = cls.on_command @wraps(old_on_command) def wrapper(self, db, message): if not message.from_admin: return RollbotResponse(message, failure=RollbotFailure.PERMISSIONS) return old_on_command(self, db, message) setattr(cls, "on_command", wrapper) return cls