From d6c61b6810eff9949d0a7d2fcc3fb43266d545b0 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Tue, 24 Mar 2020 23:07:21 +0100 Subject: Deseasonify: add generic `in_month` decorator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit See docstring for further details. This serves as a convenience wrapper around `in_month_command` and `in_month_listener` to allow a consistent API. Proposed by lemon. Co-authored-by: Leon Sandøy --- bot/decorators.py | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 19eae55b..8de2e57f 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -10,7 +10,7 @@ from weakref import WeakValueDictionary from discord import Colour, Embed from discord.ext import commands -from discord.ext.commands import CheckFailure, Context +from discord.ext.commands import CheckFailure, Command, Context from bot.constants import Client, ERROR_REPLIES, Month @@ -106,6 +106,39 @@ def in_month_command(*allowed_months: Month) -> typing.Callable: return commands.check(predicate) +def in_month(*allowed_months: Month) -> typing.Callable: + """ + Universal decorator for season-locking commands and listeners alike. + + This only serves to determine whether the decorated callable is a command, + a listener, or neither. It then delegates to either `in_month_command`, + or `in_month_listener`, or raises TypeError, respectively. + + Please note that in order for this decorator to correctly determine whether + the decorated callable is a cmd or listener, it **has** to first be turned + into one. This means that this decorator should always be placed **above** + the d.py one that registers it as either. + """ + def decorator(callable_: typing.Callable) -> typing.Callable: + # Functions decorated as commands are turned into instances of `Command` + if isinstance(callable_, Command): + logging.debug(f"Command {callable_.qualified_name} will be locked to {allowed_months}") + actual_deco = in_month_command(*allowed_months) + + # D.py will assign this attribute when `callable_` is registered as a listener + elif hasattr(callable_, "__cog_listener__"): + logging.debug(f"Listener {callable_.__qualname__} will be locked to {allowed_months}") + actual_deco = in_month_listener(*allowed_months) + + # Otherwise we're unsure exactly what has been decorated + # This happens before the bot starts, so let's just raise + else: + raise TypeError(f"Decorated object {callable_} is neither a command nor a listener") + + return actual_deco(callable_) + return decorator + + def with_role(*role_ids: int) -> typing.Callable: """Check to see whether the invoking user has any of the roles specified in role_ids.""" async def predicate(ctx: Context) -> bool: -- cgit v1.2.3