From 5fd9b74cbfe8686a1362040f81731a6ef9766970 Mon Sep 17 00:00:00 2001 From: Patchwork Collective <226386-aster.codes@users.noreply.gitlab.com> Date: Thu, 2 Sep 2021 22:23:32 -0400 Subject: [PATCH 1/2] Added simple base for a `has_any_role` decorator check. --- tanjun/checks.py | 61 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/tanjun/checks.py b/tanjun/checks.py index 759f2a290..4f901408c 100644 --- a/tanjun/checks.py +++ b/tanjun/checks.py @@ -53,6 +53,7 @@ import time import typing from collections import abc as collections +from functools import partial import hikari @@ -787,3 +788,63 @@ 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) + + +def with_any_role_check( + roles: list[str | int] = list(), + *, + error_message: str | None = "You do not have the required roles to use this command!", + halt_execution: bool = True, +) -> collections.Callable[[CommandT], CommandT]: + """Only let a command run if the author has a specific role. + + Parameters + ---------- + roles: list[str | int] + The author must have at least one (1) role in this list. (Role.name and Role.id are checked) + + Other Parameters + ---------------- + error_message: str | None + 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.""" + + def any_role_check( + ctx: tanjun_abc.Context, + *, + roles: list[str], + error_message: str | None = None, + halt_execution: bool = False, + ) -> bool: + + if not ctx.member: + return _handle_result(False, "You must be a member to use this!", True) + + member_roles = ctx.member.get_roles() + + def check_roles(member_role: Role, required_roles: list[str | int]) -> bool: + for check in required_roles: + if isinstance(check, int): + if member_role.id == check: + return True + if member_role.name == check: + return True + + result = any(check_roles(member_role, roles) for member_role in member_roles) + return _handle_result(result, error_message, halt_execution) + + partial_any_role_check = partial( + any_role_check, roles=roles, error_message=error_message, halt_execution=halt_execution + ) + return lambda command: command.add_check(partial_any_role_check) From 8936305607c05c866cbbed3fc57e4ee3013e5181 Mon Sep 17 00:00:00 2001 From: Patchwork Collective <226386-aster.codes@users.noreply.gitlab.com> Date: Fri, 3 Sep 2021 14:07:48 -0400 Subject: [PATCH 2/2] Changes based on initial PR feedback. - with_any_roles_check internal logic separated into classes for compatibility - Removed 3.10 only features --- tanjun/__init__.py | 1 + tanjun/checks.py | 92 ++++++++++++++++++++++++++-------------------- 2 files changed, 53 insertions(+), 40 deletions(-) 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 4f901408c..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 @@ -53,7 +54,6 @@ import time import typing from collections import abc as collections -from functools import partial import hikari @@ -790,61 +790,73 @@ def with_check(check: tanjun_abc.CheckSig, /) -> collections.Callable[[CommandT] 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[str | int] = list(), + roles: list[hikari.SnowflakeishOr[hikari.Role]] = list, *, - error_message: str | None = "You do not have the required roles to use this command!", + error_message: str = "You do not have the required roles to use this command!", halt_execution: bool = True, ) -> collections.Callable[[CommandT], CommandT]: - """Only let a command run if the author has a specific role. + """Add a generic check to a command. Parameters ---------- - roles: list[str | int] - The author must have at least one (1) role in this list. (Role.name and Role.id are checked) + 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: str | None + 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`. + + 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.""" - - def any_role_check( - ctx: tanjun_abc.Context, - *, - roles: list[str], - error_message: str | None = None, - halt_execution: bool = False, - ) -> bool: - - if not ctx.member: - return _handle_result(False, "You must be a member to use this!", True) - - member_roles = ctx.member.get_roles() - - def check_roles(member_role: Role, required_roles: list[str | int]) -> bool: - for check in required_roles: - if isinstance(check, int): - if member_role.id == check: - return True - if member_role.name == check: - return True - - result = any(check_roles(member_role, roles) for member_role in member_roles) - return _handle_result(result, error_message, halt_execution) - - partial_any_role_check = partial( - any_role_check, roles=roles, error_message=error_message, halt_execution=halt_execution - ) - return lambda command: command.add_check(partial_any_role_check) + A command decorator callback which adds the check. + """ + return lambda command: command.add_check(HasAnyRoleCheck(roles))