diff --git a/tanjun/__init__.py b/tanjun/__init__.py index 78b399cf2..b40f423b6 100644 --- a/tanjun/__init__.py +++ b/tanjun/__init__.py @@ -119,6 +119,7 @@ async def main() -> None: "with_owner_check", "with_author_permission_check", "with_own_permission_check", + "with_any_role_check" # clients.py "clients", "as_loader", diff --git a/tanjun/checks.py b/tanjun/checks.py index 759f2a290..20eecb4c1 100644 --- a/tanjun/checks.py +++ b/tanjun/checks.py @@ -44,6 +44,7 @@ "with_owner_check", "with_author_permission_check", "with_own_permission_check", + "with_any_role_check", ] import abc @@ -787,3 +788,75 @@ def with_check(check: tanjun_abc.CheckSig, /) -> collections.Callable[[CommandT] A command decorator callback which adds the check. """ return lambda command: command.add_check(check) + + +class HasAnyRoleCheck: + __slots__ = ( + "_halt_execution", + "_error_message", + "required_roles", + ) + + def __init__( + self, + roles: list[hikari.SnowflakeishOr[hikari.Role]] = list, + *, + error_message: str = "You do not have the required roles to use this command!", + halt_execution: bool = True, + ) -> None: + self._halt_execution = halt_execution + self._error_message = error_message + self.required_roles = roles + + async def __call__(self, ctx: tanjun_abc.Context, /) -> bool: + + if not ctx.member: + return _handle_result(False, "You must be a member to use this!", True) + + member_roles = ctx.member.get_roles() + + result = any(self.check_roles(member_role) for member_role in member_roles) + return _handle_result(result, self._error_message, self._halt_execution) + + def check_roles(self, member_role: hikari.Role) -> bool: + for check in self.required_roles: + if isinstance(check, int): + if member_role.id == check: + return True + if member_role.name == check: + return True + return False + + +def with_any_role_check( + roles: list[hikari.SnowflakeishOr[hikari.Role]] = list, + *, + error_message: str = "You do not have the required roles to use this command!", + halt_execution: bool = True, +) -> collections.Callable[[CommandT], CommandT]: + """Add a generic check to a command. + + Parameters + ---------- + roles : list[hikari.SnowflakeishOr[hikari.Role]] + The author must have at least one (1) role in this list. (Role.name and Role.id are checked). + + Other Parameters + ---------------- + error_message : Optional[str] + The error message raised if the member does not have a required role. + + Defaults to 'You do not have the required roles to use this command!' + + halt_execution : bool + Whether this check should raise `tanjun.errors.HaltExecution` to end the execution search + when it fails instead of returning `False`. + + Defaults to `False`. + + Returns + ------- + collections.abc.Callable[[CommandT], CommandT] + A command decorator callback which adds the check. + """ + return lambda command: command.add_check(HasAnyRoleCheck(roles))