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 len(message.arg_list()) < n:
                failure = RollbotFailure.INVALID_ARGUMENTS.with_reason(alert_response or f"{cls.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 len(message.arg_list()) < n:
                failure = RollbotFailure.INVALID_ARGUMENTS.with_reason(alert_response or f"{cls.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