rollbot3/lib/rollbot/plugins/decorators/requires.py

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