From 44346a9c6b08b6b30ce1ece5020ee2725218d5c9 Mon Sep 17 00:00:00 2001 From: scragly <29337040+scragly@users.noreply.github.com> Date: Fri, 6 Dec 2019 04:04:38 +1000 Subject: Fix inaccurate annotations for custom checks. --- bot/decorators.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 58f67a15..d0371df4 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -20,7 +20,7 @@ class InChannelCheckFailure(CheckFailure): pass -def with_role(*role_ids: int) -> bool: +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: if not ctx.guild: # Return False in a DM @@ -43,7 +43,7 @@ def with_role(*role_ids: int) -> bool: return commands.check(predicate) -def without_role(*role_ids: int) -> bool: +def without_role(*role_ids: int) -> typing.Callable: """Check whether the invoking user does not have all of the roles specified in role_ids.""" async def predicate(ctx: Context) -> bool: if not ctx.guild: # Return False in a DM -- cgit v1.2.3 From 428e91667f2f4b87cb1dfc811b44787d0c182e43 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Sat, 22 Feb 2020 20:23:43 +0100 Subject: Implement in_month command check Commands decorated with in_month can only be used in one of the allowed months. --- bot/decorators.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index d0371df4..8a1f00ee 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -2,6 +2,7 @@ import logging import random import typing from asyncio import Lock +from datetime import datetime from functools import wraps from weakref import WeakValueDictionary @@ -9,7 +10,7 @@ from discord import Colour, Embed from discord.ext import commands from discord.ext.commands import CheckFailure, Context -from bot.constants import ERROR_REPLIES +from bot.constants import ERROR_REPLIES, Month log = logging.getLogger(__name__) @@ -20,6 +21,24 @@ class InChannelCheckFailure(CheckFailure): pass +def in_month(*allowed_months: Month) -> typing.Callable: + """ + Check whether the command was invoked in one of `enabled_months`. + + Uses the current UTC month at the time of running the predicate. + """ + async def predicate(ctx: Context) -> bool: + current_month = datetime.utcnow().month + can_run = current_month in allowed_months + + log.debug( + f"Command '{ctx.command}' is locked to months {allowed_months}. " + f"Invoking it in month {current_month} is {'allowed' if can_run else 'disallowed'}." + ) + return can_run + return commands.check(predicate) + + 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 From 6e8b239f407bf079e35a03df41d55c042b026406 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Sun, 15 Mar 2020 14:32:12 +0100 Subject: Deseasonify: improve `in_month` command check Raise a custom exception if the command fails. This is then handled in the error handler, and the user will be informed of which months allow the invoked command. --- bot/decorators.py | 15 +++++++++++++-- bot/seasons/evergreen/error_handler.py | 4 ++-- 2 files changed, 15 insertions(+), 4 deletions(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 8a1f00ee..f031d404 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -21,6 +21,12 @@ class InChannelCheckFailure(CheckFailure): pass +class InMonthCheckFailure(CheckFailure): + """Check failure for when a command is invoked outside of its allowed month.""" + + pass + + def in_month(*allowed_months: Month) -> typing.Callable: """ Check whether the command was invoked in one of `enabled_months`. @@ -31,11 +37,16 @@ def in_month(*allowed_months: Month) -> typing.Callable: current_month = datetime.utcnow().month can_run = current_month in allowed_months + human_months = ", ".join(m.name for m in allowed_months) log.debug( - f"Command '{ctx.command}' is locked to months {allowed_months}. " + f"Command '{ctx.command}' is locked to months {human_months}. " f"Invoking it in month {current_month} is {'allowed' if can_run else 'disallowed'}." ) - return can_run + if can_run: + return True + else: + raise InMonthCheckFailure(f"Command can only be used in {human_months}") + return commands.check(predicate) diff --git a/bot/seasons/evergreen/error_handler.py b/bot/seasons/evergreen/error_handler.py index 0d8bb0bb..ba6ca5ec 100644 --- a/bot/seasons/evergreen/error_handler.py +++ b/bot/seasons/evergreen/error_handler.py @@ -7,7 +7,7 @@ from discord import Embed, Message from discord.ext import commands from bot.constants import Colours, ERROR_REPLIES, NEGATIVE_REPLIES -from bot.decorators import InChannelCheckFailure +from bot.decorators import InChannelCheckFailure, InMonthCheckFailure log = logging.getLogger(__name__) @@ -55,7 +55,7 @@ class CommandErrorHandler(commands.Cog): if isinstance(error, commands.CommandNotFound): return - if isinstance(error, InChannelCheckFailure): + if isinstance(error, (InChannelCheckFailure, InMonthCheckFailure)): await ctx.send(embed=self.error_embed(str(error), NEGATIVE_REPLIES), delete_after=7.5) return -- cgit v1.2.3 From 066436def1471ba22f1d67958e17f7d52a1ca80d Mon Sep 17 00:00:00 2001 From: kwzrd Date: Sun, 15 Mar 2020 16:33:27 +0100 Subject: Deseasonify: add listener decorator for season locking A guarded listener will abort if the triggering event happens outside of `allowed_months`. This provides a convenient way of season-locking listeners without having to write guards directly within their bodies. --- bot/decorators.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index f031d404..2520679a 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -1,3 +1,4 @@ +import functools import logging import random import typing @@ -27,6 +28,27 @@ class InMonthCheckFailure(CheckFailure): pass +def in_month_listener(*allowed_months: Month) -> typing.Callable: + """ + Shield a listener from being invoked outside of `allowed_months`. + + The check is performed against current UTC month. + """ + def decorator(listener: typing.Callable) -> typing.Callable: + @functools.wraps(listener) + async def guarded_listener(*args, **kwargs) -> None: + """Wrapped listener will abort if not in allowed month.""" + current_month = Month(datetime.utcnow().month) + + if current_month in allowed_months: + # Propagate return value although it should always be None + return await listener(*args, **kwargs) + else: + log.debug(f"Guarded {listener.__qualname__} from invoking in {current_month.name}") + return guarded_listener + return decorator + + def in_month(*allowed_months: Month) -> typing.Callable: """ Check whether the command was invoked in one of `enabled_months`. -- cgit v1.2.3 From f15732dfc630e40732273cf7bb935d8d733f19d5 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Sun, 15 Mar 2020 18:33:19 +0100 Subject: Deseasonify: add convenience decorator for seasonal tasks See docstring for implementation details. --- bot/decorators.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 2520679a..874c811b 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -1,3 +1,4 @@ +import asyncio import functools import logging import random @@ -13,6 +14,8 @@ from discord.ext.commands import CheckFailure, Context from bot.constants import ERROR_REPLIES, Month +ONE_DAY = 24 * 60 * 60 + log = logging.getLogger(__name__) @@ -28,6 +31,42 @@ class InMonthCheckFailure(CheckFailure): pass +def seasonal_task(*allowed_months: Month, sleep_time: float = ONE_DAY) -> typing.Callable: + """ + Perform the decorated method periodically in `allowed_months`. + + This provides a convenience wrapper to avoid code repetition where some task shall + perform an operation repeatedly in a constant interval, but only in specific months. + + The decorated function will be called once every `sleep_time` seconds while + the current UTC month is in `allowed_months`. Sleep time defaults to 24 hours. + """ + def decorator(task_body: typing.Callable) -> typing.Callable: + @functools.wraps(task_body) + async def decorated_task(self: commands.Cog, *args, **kwargs) -> None: + """ + Call `task_body` once every `sleep_time` seconds in `allowed_months`. + + We assume `self` to be a Cog subclass instance carrying a `bot` attr. + As some tasks may rely on the client's cache to be ready, we delegate + to the bot to wait until it's ready. + """ + await self.bot.wait_until_ready() + log.info(f"Starting seasonal task {task_body.__qualname__} ({allowed_months})") + + while True: + current_month = Month(datetime.utcnow().month) + + if current_month in allowed_months: + await task_body(self, *args, **kwargs) + else: + log.debug(f"Seasonal task {task_body.__qualname__} sleeps in {current_month.name}") + + await asyncio.sleep(sleep_time) + return decorated_task + return decorator + + def in_month_listener(*allowed_months: Month) -> typing.Callable: """ Shield a listener from being invoked outside of `allowed_months`. -- cgit v1.2.3 From f686cbacc12e06f3bb2f6586994a57b7410f7df4 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Sun, 15 Mar 2020 20:45:10 +0100 Subject: Deseasonify: remove fragile attr dependency in `seasonal_task` Importing the bot instance will allow us to safely access the `wait_until_ready` method without having to make fragile assumptions about the arguments passed to the decorated method. Although still not perfect, this feels a lot cleaner and safer than the previous approach. --- bot/decorators.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 874c811b..3ce3d8c4 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -12,6 +12,7 @@ from discord import Colour, Embed from discord.ext import commands from discord.ext.commands import CheckFailure, Context +from bot.bot import bot from bot.constants import ERROR_REPLIES, Month ONE_DAY = 24 * 60 * 60 @@ -43,22 +44,20 @@ def seasonal_task(*allowed_months: Month, sleep_time: float = ONE_DAY) -> typing """ def decorator(task_body: typing.Callable) -> typing.Callable: @functools.wraps(task_body) - async def decorated_task(self: commands.Cog, *args, **kwargs) -> None: + async def decorated_task(*args, **kwargs) -> None: """ Call `task_body` once every `sleep_time` seconds in `allowed_months`. - We assume `self` to be a Cog subclass instance carrying a `bot` attr. - As some tasks may rely on the client's cache to be ready, we delegate - to the bot to wait until it's ready. + Wait for bot to be ready before calling `task_body` for the first time. """ - await self.bot.wait_until_ready() + await bot.wait_until_ready() log.info(f"Starting seasonal task {task_body.__qualname__} ({allowed_months})") while True: current_month = Month(datetime.utcnow().month) if current_month in allowed_months: - await task_body(self, *args, **kwargs) + await task_body(*args, **kwargs) else: log.debug(f"Seasonal task {task_body.__qualname__} sleeps in {current_month.name}") -- cgit v1.2.3 From f9cef6ae66e49c566d416a3b2bad7d34c3704a80 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Sun, 15 Mar 2020 20:47:02 +0100 Subject: Deseasonify: relax type annotation Asyncio's sleep will accept both, and we default to an int, so might as well not break our own promise. --- bot/decorators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 3ce3d8c4..74976cd6 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -32,7 +32,7 @@ class InMonthCheckFailure(CheckFailure): pass -def seasonal_task(*allowed_months: Month, sleep_time: float = ONE_DAY) -> typing.Callable: +def seasonal_task(*allowed_months: Month, sleep_time: typing.Union[float, int] = ONE_DAY) -> typing.Callable: """ Perform the decorated method periodically in `allowed_months`. -- cgit v1.2.3 From 79376abcb7d305484bf5283f95d3a32641aeb6d5 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Mon, 23 Mar 2020 19:31:23 +0100 Subject: Deseasonify: abdicate responsibility to wait until bot is ready The wrapper will no longer wait for the bot to be ready before it calls the wrapped function for the first time. If the function requires the bot's cache to be ready, it is responsible for awaiting the method itself. This removes the need to acquire a reference to the bot instance inside the decorator, which seems to be difficult to do cleanly. As Mark adds, this may in fact be safer as the bot may temporarily disconnect while the task is active, and awaiting the bot's ready status every time will prevent issues in such a situation. Co-authored-by: MarkKoz --- bot/decorators.py | 10 +++------- bot/seasons/easter/egg_facts.py | 2 ++ bot/seasons/pride/pride_facts.py | 2 ++ 3 files changed, 7 insertions(+), 7 deletions(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 74976cd6..400f1bbb 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -12,7 +12,6 @@ from discord import Colour, Embed from discord.ext import commands from discord.ext.commands import CheckFailure, Context -from bot.bot import bot from bot.constants import ERROR_REPLIES, Month ONE_DAY = 24 * 60 * 60 @@ -41,16 +40,13 @@ def seasonal_task(*allowed_months: Month, sleep_time: typing.Union[float, int] = The decorated function will be called once every `sleep_time` seconds while the current UTC month is in `allowed_months`. Sleep time defaults to 24 hours. + + The wrapped task is responsible for waiting for the bot to be ready, if necessary. """ def decorator(task_body: typing.Callable) -> typing.Callable: @functools.wraps(task_body) async def decorated_task(*args, **kwargs) -> None: - """ - Call `task_body` once every `sleep_time` seconds in `allowed_months`. - - Wait for bot to be ready before calling `task_body` for the first time. - """ - await bot.wait_until_ready() + """Call `task_body` once every `sleep_time` seconds in `allowed_months`.""" log.info(f"Starting seasonal task {task_body.__qualname__} ({allowed_months})") while True: diff --git a/bot/seasons/easter/egg_facts.py b/bot/seasons/easter/egg_facts.py index f61f9da4..d20df3de 100644 --- a/bot/seasons/easter/egg_facts.py +++ b/bot/seasons/easter/egg_facts.py @@ -35,6 +35,8 @@ class EasterFacts(commands.Cog): @seasonal_task(Month.april) async def send_egg_fact_daily(self) -> None: """A background task that sends an easter egg fact in the event channel everyday.""" + await self.bot.wait_until_ready() + channel = self.bot.get_channel(Channels.seasonalbot_commands) await channel.send(embed=self.make_embed()) diff --git a/bot/seasons/pride/pride_facts.py b/bot/seasons/pride/pride_facts.py index 417a49a6..1a02eaaa 100644 --- a/bot/seasons/pride/pride_facts.py +++ b/bot/seasons/pride/pride_facts.py @@ -35,6 +35,8 @@ class PrideFacts(commands.Cog): @seasonal_task(Month.june) async def send_pride_fact_daily(self) -> None: """Background task to post the daily pride fact every day.""" + await self.bot.wait_until_ready() + channel = self.bot.get_channel(Channels.seasonalbot_commands) await self.send_select_fact(channel, datetime.utcnow()) -- cgit v1.2.3 From 1a326904f1bfead7c2b25839910fbeb4c70d84fb Mon Sep 17 00:00:00 2001 From: kwzrd Date: Mon, 23 Mar 2020 20:17:12 +0100 Subject: Deseasonify: add `mock_in_debug` decorator This should be very useful for testing. See docstring. --- bot/decorators.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 400f1bbb..efd43da0 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -12,7 +12,7 @@ from discord import Colour, Embed from discord.ext import commands from discord.ext.commands import CheckFailure, Context -from bot.constants import ERROR_REPLIES, Month +from bot.constants import Client, ERROR_REPLIES, Month ONE_DAY = 24 * 60 * 60 @@ -261,3 +261,23 @@ def locked() -> typing.Union[typing.Callable, None]: return await func(self, ctx, *args, **kwargs) return inner return wrap + + +def mock_in_debug(return_value: typing.Any) -> typing.Callable: + """ + Short-circuit function execution if in debug mode and return `return_value`. + + The original function name, and the incoming args and kwargs are DEBUG level logged + upon each call. This is useful for expensive operations, i.e. media asset uploads + that are prone to rate-limits but need to be tested extensively. + """ + def decorator(func: typing.Callable) -> typing.Callable: + @functools.wraps(func) + async def wrapped(*args, **kwargs) -> typing.Any: + """Short-circuit and log if in debug mode.""" + if Client.debug: + log.debug(f"Function {func.__name__} called with args: {args}, kwargs: {kwargs}") + return return_value + return await func(*args, **kwargs) + return wrapped + return decorator -- cgit v1.2.3 From 28e4a23dd50a047dce76818a410d4437c0d0c443 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Mon, 23 Mar 2020 22:41:06 +0100 Subject: Deseasonify: rename `in_month` decorator Indicate that the decorator shall only be applied to commands. The `in_month` name will be used for a universal decorator that can season-lock both listeners and commands. --- bot/decorators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index efd43da0..19eae55b 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -83,7 +83,7 @@ def in_month_listener(*allowed_months: Month) -> typing.Callable: return decorator -def in_month(*allowed_months: Month) -> typing.Callable: +def in_month_command(*allowed_months: Month) -> typing.Callable: """ Check whether the command was invoked in one of `enabled_months`. -- cgit v1.2.3 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 From 56eb5c4108976ed125984eb67ed0ef5d484813d4 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Tue, 24 Mar 2020 23:12:07 +0100 Subject: Deseasonify: use 't-dot' notation for type annotations The module is full of complicated annotations, and the full `typing` takes up annoyingly much visual space. --- bot/decorators.py | 38 +++++++++++++++++++------------------- 1 file changed, 19 insertions(+), 19 deletions(-) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 8de2e57f..0a1f77c8 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -2,7 +2,7 @@ import asyncio import functools import logging import random -import typing +import typing as t from asyncio import Lock from datetime import datetime from functools import wraps @@ -31,7 +31,7 @@ class InMonthCheckFailure(CheckFailure): pass -def seasonal_task(*allowed_months: Month, sleep_time: typing.Union[float, int] = ONE_DAY) -> typing.Callable: +def seasonal_task(*allowed_months: Month, sleep_time: t.Union[float, int] = ONE_DAY) -> t.Callable: """ Perform the decorated method periodically in `allowed_months`. @@ -43,7 +43,7 @@ def seasonal_task(*allowed_months: Month, sleep_time: typing.Union[float, int] = The wrapped task is responsible for waiting for the bot to be ready, if necessary. """ - def decorator(task_body: typing.Callable) -> typing.Callable: + def decorator(task_body: t.Callable) -> t.Callable: @functools.wraps(task_body) async def decorated_task(*args, **kwargs) -> None: """Call `task_body` once every `sleep_time` seconds in `allowed_months`.""" @@ -62,13 +62,13 @@ def seasonal_task(*allowed_months: Month, sleep_time: typing.Union[float, int] = return decorator -def in_month_listener(*allowed_months: Month) -> typing.Callable: +def in_month_listener(*allowed_months: Month) -> t.Callable: """ Shield a listener from being invoked outside of `allowed_months`. The check is performed against current UTC month. """ - def decorator(listener: typing.Callable) -> typing.Callable: + def decorator(listener: t.Callable) -> t.Callable: @functools.wraps(listener) async def guarded_listener(*args, **kwargs) -> None: """Wrapped listener will abort if not in allowed month.""" @@ -83,7 +83,7 @@ def in_month_listener(*allowed_months: Month) -> typing.Callable: return decorator -def in_month_command(*allowed_months: Month) -> typing.Callable: +def in_month_command(*allowed_months: Month) -> t.Callable: """ Check whether the command was invoked in one of `enabled_months`. @@ -106,7 +106,7 @@ def in_month_command(*allowed_months: Month) -> typing.Callable: return commands.check(predicate) -def in_month(*allowed_months: Month) -> typing.Callable: +def in_month(*allowed_months: Month) -> t.Callable: """ Universal decorator for season-locking commands and listeners alike. @@ -119,7 +119,7 @@ def in_month(*allowed_months: Month) -> typing.Callable: 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: + def decorator(callable_: t.Callable) -> t.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}") @@ -139,7 +139,7 @@ def in_month(*allowed_months: Month) -> typing.Callable: return decorator -def with_role(*role_ids: int) -> typing.Callable: +def with_role(*role_ids: int) -> t.Callable: """Check to see whether the invoking user has any of the roles specified in role_ids.""" async def predicate(ctx: Context) -> bool: if not ctx.guild: # Return False in a DM @@ -162,7 +162,7 @@ def with_role(*role_ids: int) -> typing.Callable: return commands.check(predicate) -def without_role(*role_ids: int) -> typing.Callable: +def without_role(*role_ids: int) -> t.Callable: """Check whether the invoking user does not have all of the roles specified in role_ids.""" async def predicate(ctx: Context) -> bool: if not ctx.guild: # Return False in a DM @@ -182,7 +182,7 @@ def without_role(*role_ids: int) -> typing.Callable: return commands.check(predicate) -def in_channel_check(*channels: int, bypass_roles: typing.Container[int] = None) -> typing.Callable[[Context], bool]: +def in_channel_check(*channels: int, bypass_roles: t.Container[int] = None) -> t.Callable[[Context], bool]: """ Checks that the message is in a whitelisted channel or optionally has a bypass role. @@ -248,7 +248,7 @@ def in_channel_check(*channels: int, bypass_roles: typing.Container[int] = None) in_channel = commands.check(in_channel_check) -def override_in_channel(channels: typing.Tuple[int] = None) -> typing.Callable: +def override_in_channel(channels: t.Tuple[int] = None) -> t.Callable: """ Set command callback attribute for detection in `in_channel_check`. @@ -256,14 +256,14 @@ def override_in_channel(channels: typing.Tuple[int] = None) -> typing.Callable: This decorator has to go before (below) below the `command` decorator. """ - def inner(func: typing.Callable) -> typing.Callable: + def inner(func: t.Callable) -> t.Callable: func.in_channel_override = channels return func return inner -def locked() -> typing.Union[typing.Callable, None]: +def locked() -> t.Union[t.Callable, None]: """ Allows the user to only run one instance of the decorated command at a time. @@ -271,11 +271,11 @@ def locked() -> typing.Union[typing.Callable, None]: This decorator has to go before (below) the `command` decorator. """ - def wrap(func: typing.Callable) -> typing.Union[typing.Callable, None]: + def wrap(func: t.Callable) -> t.Union[t.Callable, None]: func.__locks = WeakValueDictionary() @wraps(func) - async def inner(self: typing.Callable, ctx: Context, *args, **kwargs) -> typing.Union[typing.Callable, None]: + async def inner(self: t.Callable, ctx: Context, *args, **kwargs) -> t.Union[t.Callable, None]: lock = func.__locks.setdefault(ctx.author.id, Lock()) if lock.locked(): embed = Embed() @@ -296,7 +296,7 @@ def locked() -> typing.Union[typing.Callable, None]: return wrap -def mock_in_debug(return_value: typing.Any) -> typing.Callable: +def mock_in_debug(return_value: t.Any) -> t.Callable: """ Short-circuit function execution if in debug mode and return `return_value`. @@ -304,9 +304,9 @@ def mock_in_debug(return_value: typing.Any) -> typing.Callable: upon each call. This is useful for expensive operations, i.e. media asset uploads that are prone to rate-limits but need to be tested extensively. """ - def decorator(func: typing.Callable) -> typing.Callable: + def decorator(func: t.Callable) -> t.Callable: @functools.wraps(func) - async def wrapped(*args, **kwargs) -> typing.Any: + async def wrapped(*args, **kwargs) -> t.Any: """Short-circuit and log if in debug mode.""" if Client.debug: log.debug(f"Function {func.__name__} called with args: {args}, kwargs: {kwargs}") -- cgit v1.2.3 From c04103b99e0970ac1d4088f1230b3484c71dbee9 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Thu, 26 Mar 2020 21:17:55 +0100 Subject: Deseasonify: extend `in_month` doc --- bot/decorators.py | 5 +++++ 1 file changed, 5 insertions(+) (limited to 'bot/decorators.py') diff --git a/bot/decorators.py b/bot/decorators.py index 0a1f77c8..f85996b5 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -118,6 +118,11 @@ def in_month(*allowed_months: Month) -> t.Callable: 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. + + This will decorate groups as well, as those subclass Command. In order to lock + all subcommands of a group, its `invoke_without_command` param must **not** be + manually set to True - this causes a circumvention of the group's callback + and the seasonal check applied to it. """ def decorator(callable_: t.Callable) -> t.Callable: # Functions decorated as commands are turned into instances of `Command` -- cgit v1.2.3 From 99b23ff1500aa15b076ee2e1c5e1b5560ba2c366 Mon Sep 17 00:00:00 2001 From: kwzrd Date: Sat, 28 Mar 2020 14:51:14 +0100 Subject: Deseasonify: move decorators module under utils --- bot/__main__.py | 2 +- bot/bot.py | 2 +- bot/branding.py | 2 +- bot/decorators.py | 321 ------------------------ bot/seasons/christmas/adventofcode.py | 2 +- bot/seasons/christmas/hanukkah_embed.py | 2 +- bot/seasons/easter/egg_facts.py | 2 +- bot/seasons/evergreen/error_handler.py | 2 +- bot/seasons/evergreen/game.py | 2 +- bot/seasons/evergreen/issues.py | 2 +- bot/seasons/evergreen/snakes/snakes_cog.py | 2 +- bot/seasons/halloween/candy_collection.py | 2 +- bot/seasons/halloween/hacktober-issue-finder.py | 2 +- bot/seasons/halloween/hacktoberstats.py | 2 +- bot/seasons/halloween/spookyreact.py | 2 +- bot/seasons/pride/pride_facts.py | 2 +- bot/seasons/valentines/be_my_valentine.py | 2 +- bot/utils/decorators.py | 321 ++++++++++++++++++++++++ 18 files changed, 337 insertions(+), 337 deletions(-) delete mode 100644 bot/decorators.py create mode 100644 bot/utils/decorators.py (limited to 'bot/decorators.py') diff --git a/bot/__main__.py b/bot/__main__.py index 780c8c4d..3662a63b 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -2,8 +2,8 @@ import logging from bot.bot import bot from bot.constants import Client, STAFF_ROLES, WHITELISTED_CHANNELS -from bot.decorators import in_channel_check from bot.seasons import get_extensions +from bot.utils.decorators import in_channel_check log = logging.getLogger(__name__) diff --git a/bot/bot.py b/bot/bot.py index b47e1289..47d63de9 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -11,7 +11,7 @@ from discord import DiscordException, Embed from discord.ext import commands from bot.constants import Channels, Client -from bot.decorators import mock_in_debug +from bot.utils.decorators import mock_in_debug log = logging.getLogger(__name__) diff --git a/bot/branding.py b/bot/branding.py index 6f5e7a5c..2eb563ea 100644 --- a/bot/branding.py +++ b/bot/branding.py @@ -12,8 +12,8 @@ from discord.ext import commands from bot.bot import SeasonalBot from bot.constants import Branding, Colours, Emojis, MODERATION_ROLES, Tokens -from bot.decorators import with_role from bot.seasons import SeasonBase, get_current_season, get_season +from bot.utils.decorators import with_role from bot.utils.exceptions import BrandingError log = logging.getLogger(__name__) diff --git a/bot/decorators.py b/bot/decorators.py deleted file mode 100644 index f85996b5..00000000 --- a/bot/decorators.py +++ /dev/null @@ -1,321 +0,0 @@ -import asyncio -import functools -import logging -import random -import typing as t -from asyncio import Lock -from datetime import datetime -from functools import wraps -from weakref import WeakValueDictionary - -from discord import Colour, Embed -from discord.ext import commands -from discord.ext.commands import CheckFailure, Command, Context - -from bot.constants import Client, ERROR_REPLIES, Month - -ONE_DAY = 24 * 60 * 60 - -log = logging.getLogger(__name__) - - -class InChannelCheckFailure(CheckFailure): - """Check failure when the user runs a command in a non-whitelisted channel.""" - - pass - - -class InMonthCheckFailure(CheckFailure): - """Check failure for when a command is invoked outside of its allowed month.""" - - pass - - -def seasonal_task(*allowed_months: Month, sleep_time: t.Union[float, int] = ONE_DAY) -> t.Callable: - """ - Perform the decorated method periodically in `allowed_months`. - - This provides a convenience wrapper to avoid code repetition where some task shall - perform an operation repeatedly in a constant interval, but only in specific months. - - The decorated function will be called once every `sleep_time` seconds while - the current UTC month is in `allowed_months`. Sleep time defaults to 24 hours. - - The wrapped task is responsible for waiting for the bot to be ready, if necessary. - """ - def decorator(task_body: t.Callable) -> t.Callable: - @functools.wraps(task_body) - async def decorated_task(*args, **kwargs) -> None: - """Call `task_body` once every `sleep_time` seconds in `allowed_months`.""" - log.info(f"Starting seasonal task {task_body.__qualname__} ({allowed_months})") - - while True: - current_month = Month(datetime.utcnow().month) - - if current_month in allowed_months: - await task_body(*args, **kwargs) - else: - log.debug(f"Seasonal task {task_body.__qualname__} sleeps in {current_month.name}") - - await asyncio.sleep(sleep_time) - return decorated_task - return decorator - - -def in_month_listener(*allowed_months: Month) -> t.Callable: - """ - Shield a listener from being invoked outside of `allowed_months`. - - The check is performed against current UTC month. - """ - def decorator(listener: t.Callable) -> t.Callable: - @functools.wraps(listener) - async def guarded_listener(*args, **kwargs) -> None: - """Wrapped listener will abort if not in allowed month.""" - current_month = Month(datetime.utcnow().month) - - if current_month in allowed_months: - # Propagate return value although it should always be None - return await listener(*args, **kwargs) - else: - log.debug(f"Guarded {listener.__qualname__} from invoking in {current_month.name}") - return guarded_listener - return decorator - - -def in_month_command(*allowed_months: Month) -> t.Callable: - """ - Check whether the command was invoked in one of `enabled_months`. - - Uses the current UTC month at the time of running the predicate. - """ - async def predicate(ctx: Context) -> bool: - current_month = datetime.utcnow().month - can_run = current_month in allowed_months - - human_months = ", ".join(m.name for m in allowed_months) - log.debug( - f"Command '{ctx.command}' is locked to months {human_months}. " - f"Invoking it in month {current_month} is {'allowed' if can_run else 'disallowed'}." - ) - if can_run: - return True - else: - raise InMonthCheckFailure(f"Command can only be used in {human_months}") - - return commands.check(predicate) - - -def in_month(*allowed_months: Month) -> t.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. - - This will decorate groups as well, as those subclass Command. In order to lock - all subcommands of a group, its `invoke_without_command` param must **not** be - manually set to True - this causes a circumvention of the group's callback - and the seasonal check applied to it. - """ - def decorator(callable_: t.Callable) -> t.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) -> t.Callable: - """Check to see whether the invoking user has any of the roles specified in role_ids.""" - async def predicate(ctx: Context) -> bool: - if not ctx.guild: # Return False in a DM - log.debug( - f"{ctx.author} tried to use the '{ctx.command.name}'command from a DM. " - "This command is restricted by the with_role decorator. Rejecting request." - ) - return False - - for role in ctx.author.roles: - if role.id in role_ids: - log.debug(f"{ctx.author} has the '{role.name}' role, and passes the check.") - return True - - log.debug( - f"{ctx.author} does not have the required role to use " - f"the '{ctx.command.name}' command, so the request is rejected." - ) - return False - return commands.check(predicate) - - -def without_role(*role_ids: int) -> t.Callable: - """Check whether the invoking user does not have all of the roles specified in role_ids.""" - async def predicate(ctx: Context) -> bool: - if not ctx.guild: # Return False in a DM - log.debug( - f"{ctx.author} tried to use the '{ctx.command.name}' command from a DM. " - "This command is restricted by the without_role decorator. Rejecting request." - ) - return False - - author_roles = [role.id for role in ctx.author.roles] - check = all(role not in author_roles for role in role_ids) - log.debug( - f"{ctx.author} tried to call the '{ctx.command.name}' command. " - f"The result of the without_role check was {check}." - ) - return check - return commands.check(predicate) - - -def in_channel_check(*channels: int, bypass_roles: t.Container[int] = None) -> t.Callable[[Context], bool]: - """ - Checks that the message is in a whitelisted channel or optionally has a bypass role. - - If `in_channel_override` is present, check if it contains channels - and use them in place of the global whitelist. - """ - def predicate(ctx: Context) -> bool: - if not ctx.guild: - log.debug(f"{ctx.author} tried to use the '{ctx.command.name}' command from a DM.") - return True - if ctx.channel.id in channels: - log.debug( - f"{ctx.author} tried to call the '{ctx.command.name}' command " - f"and the command was used in a whitelisted channel." - ) - return True - - if bypass_roles and any(r.id in bypass_roles for r in ctx.author.roles): - log.debug( - f"{ctx.author} called the '{ctx.command.name}' command and " - f"had a role to bypass the in_channel check." - ) - return True - - if hasattr(ctx.command.callback, "in_channel_override"): - override = ctx.command.callback.in_channel_override - if override is None: - log.debug( - f"{ctx.author} called the '{ctx.command.name}' command " - f"and the command was whitelisted to bypass the in_channel check." - ) - return True - else: - if ctx.channel.id in override: - log.debug( - f"{ctx.author} tried to call the '{ctx.command.name}' command " - f"and the command was used in an overridden whitelisted channel." - ) - return True - - log.debug( - f"{ctx.author} tried to call the '{ctx.command.name}' command. " - f"The overridden in_channel check failed." - ) - channels_str = ', '.join(f"<#{c_id}>" for c_id in override) - raise InChannelCheckFailure( - f"Sorry, but you may only use this command within {channels_str}." - ) - - log.debug( - f"{ctx.author} tried to call the '{ctx.command.name}' command. " - f"The in_channel check failed." - ) - - channels_str = ', '.join(f"<#{c_id}>" for c_id in channels) - raise InChannelCheckFailure( - f"Sorry, but you may only use this command within {channels_str}." - ) - - return predicate - - -in_channel = commands.check(in_channel_check) - - -def override_in_channel(channels: t.Tuple[int] = None) -> t.Callable: - """ - Set command callback attribute for detection in `in_channel_check`. - - Override global whitelist if channels are specified. - - This decorator has to go before (below) below the `command` decorator. - """ - def inner(func: t.Callable) -> t.Callable: - func.in_channel_override = channels - return func - - return inner - - -def locked() -> t.Union[t.Callable, None]: - """ - Allows the user to only run one instance of the decorated command at a time. - - Subsequent calls to the command from the same author are ignored until the command has completed invocation. - - This decorator has to go before (below) the `command` decorator. - """ - def wrap(func: t.Callable) -> t.Union[t.Callable, None]: - func.__locks = WeakValueDictionary() - - @wraps(func) - async def inner(self: t.Callable, ctx: Context, *args, **kwargs) -> t.Union[t.Callable, None]: - lock = func.__locks.setdefault(ctx.author.id, Lock()) - if lock.locked(): - embed = Embed() - embed.colour = Colour.red() - - log.debug(f"User tried to invoke a locked command.") - embed.description = ( - "You're already using this command. Please wait until " - "it is done before you use it again." - ) - embed.title = random.choice(ERROR_REPLIES) - await ctx.send(embed=embed) - return - - async with func.__locks.setdefault(ctx.author.id, Lock()): - return await func(self, ctx, *args, **kwargs) - return inner - return wrap - - -def mock_in_debug(return_value: t.Any) -> t.Callable: - """ - Short-circuit function execution if in debug mode and return `return_value`. - - The original function name, and the incoming args and kwargs are DEBUG level logged - upon each call. This is useful for expensive operations, i.e. media asset uploads - that are prone to rate-limits but need to be tested extensively. - """ - def decorator(func: t.Callable) -> t.Callable: - @functools.wraps(func) - async def wrapped(*args, **kwargs) -> t.Any: - """Short-circuit and log if in debug mode.""" - if Client.debug: - log.debug(f"Function {func.__name__} called with args: {args}, kwargs: {kwargs}") - return return_value - return await func(*args, **kwargs) - return wrapped - return decorator diff --git a/bot/seasons/christmas/adventofcode.py b/bot/seasons/christmas/adventofcode.py index e6100056..f7590e04 100644 --- a/bot/seasons/christmas/adventofcode.py +++ b/bot/seasons/christmas/adventofcode.py @@ -14,8 +14,8 @@ from discord.ext import commands from pytz import timezone from bot.constants import AdventOfCode as AocConfig, Channels, Colours, Emojis, Month, Tokens, WHITELISTED_CHANNELS -from bot.decorators import in_month, override_in_channel from bot.utils import unlocked_role +from bot.utils.decorators import in_month, override_in_channel log = logging.getLogger(__name__) diff --git a/bot/seasons/christmas/hanukkah_embed.py b/bot/seasons/christmas/hanukkah_embed.py index e73a33ad..62efd04e 100644 --- a/bot/seasons/christmas/hanukkah_embed.py +++ b/bot/seasons/christmas/hanukkah_embed.py @@ -6,7 +6,7 @@ from discord import Embed from discord.ext import commands from bot.constants import Colours, Month -from bot.decorators import in_month +from bot.utils.decorators import in_month log = logging.getLogger(__name__) diff --git a/bot/seasons/easter/egg_facts.py b/bot/seasons/easter/egg_facts.py index d20df3de..99a80b28 100644 --- a/bot/seasons/easter/egg_facts.py +++ b/bot/seasons/easter/egg_facts.py @@ -7,7 +7,7 @@ import discord from discord.ext import commands from bot.constants import Channels, Colours, Month -from bot.decorators import seasonal_task +from bot.utils.decorators import seasonal_task log = logging.getLogger(__name__) diff --git a/bot/seasons/evergreen/error_handler.py b/bot/seasons/evergreen/error_handler.py index d2accbd1..d268dab1 100644 --- a/bot/seasons/evergreen/error_handler.py +++ b/bot/seasons/evergreen/error_handler.py @@ -7,7 +7,7 @@ from discord import Embed, Message from discord.ext import commands from bot.constants import Colours, ERROR_REPLIES, NEGATIVE_REPLIES -from bot.decorators import InChannelCheckFailure, InMonthCheckFailure +from bot.utils.decorators import InChannelCheckFailure, InMonthCheckFailure from bot.utils.exceptions import BrandingError log = logging.getLogger(__name__) diff --git a/bot/seasons/evergreen/game.py b/bot/seasons/evergreen/game.py index ace77e9d..d43b1ad6 100644 --- a/bot/seasons/evergreen/game.py +++ b/bot/seasons/evergreen/game.py @@ -12,7 +12,7 @@ from discord.ext.commands import Cog, Context, group from bot.bot import SeasonalBot from bot.constants import STAFF_ROLES, Tokens -from bot.decorators import with_role +from bot.utils.decorators import with_role from bot.utils.pagination import ImagePaginator, LinePaginator # Base URL of IGDB API diff --git a/bot/seasons/evergreen/issues.py b/bot/seasons/evergreen/issues.py index fba5b174..fb18b62a 100644 --- a/bot/seasons/evergreen/issues.py +++ b/bot/seasons/evergreen/issues.py @@ -4,7 +4,7 @@ import discord from discord.ext import commands from bot.constants import Colours, Emojis, WHITELISTED_CHANNELS -from bot.decorators import override_in_channel +from bot.utils.decorators import override_in_channel log = logging.getLogger(__name__) diff --git a/bot/seasons/evergreen/snakes/snakes_cog.py b/bot/seasons/evergreen/snakes/snakes_cog.py index 09f5e250..e5a03a20 100644 --- a/bot/seasons/evergreen/snakes/snakes_cog.py +++ b/bot/seasons/evergreen/snakes/snakes_cog.py @@ -18,9 +18,9 @@ from discord import Colour, Embed, File, Member, Message, Reaction from discord.ext.commands import BadArgument, Bot, Cog, CommandError, Context, bot_has_permissions, group from bot.constants import ERROR_REPLIES, Tokens -from bot.decorators import locked from bot.seasons.evergreen.snakes import utils from bot.seasons.evergreen.snakes.converter import Snake +from bot.utils.decorators import locked log = logging.getLogger(__name__) diff --git a/bot/seasons/halloween/candy_collection.py b/bot/seasons/halloween/candy_collection.py index 967a62aa..3f2b895e 100644 --- a/bot/seasons/halloween/candy_collection.py +++ b/bot/seasons/halloween/candy_collection.py @@ -9,7 +9,7 @@ import discord from discord.ext import commands from bot.constants import Channels, Month -from bot.decorators import in_month +from bot.utils.decorators import in_month log = logging.getLogger(__name__) diff --git a/bot/seasons/halloween/hacktober-issue-finder.py b/bot/seasons/halloween/hacktober-issue-finder.py index e90796f1..f15a665a 100644 --- a/bot/seasons/halloween/hacktober-issue-finder.py +++ b/bot/seasons/halloween/hacktober-issue-finder.py @@ -8,7 +8,7 @@ import discord from discord.ext import commands from bot.constants import Month -from bot.decorators import in_month +from bot.utils.decorators import in_month log = logging.getLogger(__name__) diff --git a/bot/seasons/halloween/hacktoberstats.py b/bot/seasons/halloween/hacktoberstats.py index 3b1444ab..5dfa2f51 100644 --- a/bot/seasons/halloween/hacktoberstats.py +++ b/bot/seasons/halloween/hacktoberstats.py @@ -11,7 +11,7 @@ import discord from discord.ext import commands from bot.constants import Channels, Month, WHITELISTED_CHANNELS -from bot.decorators import in_month, override_in_channel +from bot.utils.decorators import in_month, override_in_channel from bot.utils.persist import make_persistent log = logging.getLogger(__name__) diff --git a/bot/seasons/halloween/spookyreact.py b/bot/seasons/halloween/spookyreact.py index 37f42a86..16f18019 100644 --- a/bot/seasons/halloween/spookyreact.py +++ b/bot/seasons/halloween/spookyreact.py @@ -5,7 +5,7 @@ import discord from discord.ext.commands import Bot, Cog from bot.constants import Month -from bot.decorators import in_month +from bot.utils.decorators import in_month log = logging.getLogger(__name__) diff --git a/bot/seasons/pride/pride_facts.py b/bot/seasons/pride/pride_facts.py index 1a02eaaa..2db8f5c2 100644 --- a/bot/seasons/pride/pride_facts.py +++ b/bot/seasons/pride/pride_facts.py @@ -10,7 +10,7 @@ import discord from discord.ext import commands from bot.constants import Channels, Colours, Month -from bot.decorators import seasonal_task +from bot.utils.decorators import seasonal_task log = logging.getLogger(__name__) diff --git a/bot/seasons/valentines/be_my_valentine.py b/bot/seasons/valentines/be_my_valentine.py index 67d8796a..1e883d21 100644 --- a/bot/seasons/valentines/be_my_valentine.py +++ b/bot/seasons/valentines/be_my_valentine.py @@ -9,7 +9,7 @@ from discord.ext import commands from discord.ext.commands.cooldowns import BucketType from bot.constants import Channels, Client, Colours, Lovefest, Month -from bot.decorators import in_month +from bot.utils.decorators import in_month log = logging.getLogger(__name__) diff --git a/bot/utils/decorators.py b/bot/utils/decorators.py new file mode 100644 index 00000000..f85996b5 --- /dev/null +++ b/bot/utils/decorators.py @@ -0,0 +1,321 @@ +import asyncio +import functools +import logging +import random +import typing as t +from asyncio import Lock +from datetime import datetime +from functools import wraps +from weakref import WeakValueDictionary + +from discord import Colour, Embed +from discord.ext import commands +from discord.ext.commands import CheckFailure, Command, Context + +from bot.constants import Client, ERROR_REPLIES, Month + +ONE_DAY = 24 * 60 * 60 + +log = logging.getLogger(__name__) + + +class InChannelCheckFailure(CheckFailure): + """Check failure when the user runs a command in a non-whitelisted channel.""" + + pass + + +class InMonthCheckFailure(CheckFailure): + """Check failure for when a command is invoked outside of its allowed month.""" + + pass + + +def seasonal_task(*allowed_months: Month, sleep_time: t.Union[float, int] = ONE_DAY) -> t.Callable: + """ + Perform the decorated method periodically in `allowed_months`. + + This provides a convenience wrapper to avoid code repetition where some task shall + perform an operation repeatedly in a constant interval, but only in specific months. + + The decorated function will be called once every `sleep_time` seconds while + the current UTC month is in `allowed_months`. Sleep time defaults to 24 hours. + + The wrapped task is responsible for waiting for the bot to be ready, if necessary. + """ + def decorator(task_body: t.Callable) -> t.Callable: + @functools.wraps(task_body) + async def decorated_task(*args, **kwargs) -> None: + """Call `task_body` once every `sleep_time` seconds in `allowed_months`.""" + log.info(f"Starting seasonal task {task_body.__qualname__} ({allowed_months})") + + while True: + current_month = Month(datetime.utcnow().month) + + if current_month in allowed_months: + await task_body(*args, **kwargs) + else: + log.debug(f"Seasonal task {task_body.__qualname__} sleeps in {current_month.name}") + + await asyncio.sleep(sleep_time) + return decorated_task + return decorator + + +def in_month_listener(*allowed_months: Month) -> t.Callable: + """ + Shield a listener from being invoked outside of `allowed_months`. + + The check is performed against current UTC month. + """ + def decorator(listener: t.Callable) -> t.Callable: + @functools.wraps(listener) + async def guarded_listener(*args, **kwargs) -> None: + """Wrapped listener will abort if not in allowed month.""" + current_month = Month(datetime.utcnow().month) + + if current_month in allowed_months: + # Propagate return value although it should always be None + return await listener(*args, **kwargs) + else: + log.debug(f"Guarded {listener.__qualname__} from invoking in {current_month.name}") + return guarded_listener + return decorator + + +def in_month_command(*allowed_months: Month) -> t.Callable: + """ + Check whether the command was invoked in one of `enabled_months`. + + Uses the current UTC month at the time of running the predicate. + """ + async def predicate(ctx: Context) -> bool: + current_month = datetime.utcnow().month + can_run = current_month in allowed_months + + human_months = ", ".join(m.name for m in allowed_months) + log.debug( + f"Command '{ctx.command}' is locked to months {human_months}. " + f"Invoking it in month {current_month} is {'allowed' if can_run else 'disallowed'}." + ) + if can_run: + return True + else: + raise InMonthCheckFailure(f"Command can only be used in {human_months}") + + return commands.check(predicate) + + +def in_month(*allowed_months: Month) -> t.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. + + This will decorate groups as well, as those subclass Command. In order to lock + all subcommands of a group, its `invoke_without_command` param must **not** be + manually set to True - this causes a circumvention of the group's callback + and the seasonal check applied to it. + """ + def decorator(callable_: t.Callable) -> t.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) -> t.Callable: + """Check to see whether the invoking user has any of the roles specified in role_ids.""" + async def predicate(ctx: Context) -> bool: + if not ctx.guild: # Return False in a DM + log.debug( + f"{ctx.author} tried to use the '{ctx.command.name}'command from a DM. " + "This command is restricted by the with_role decorator. Rejecting request." + ) + return False + + for role in ctx.author.roles: + if role.id in role_ids: + log.debug(f"{ctx.author} has the '{role.name}' role, and passes the check.") + return True + + log.debug( + f"{ctx.author} does not have the required role to use " + f"the '{ctx.command.name}' command, so the request is rejected." + ) + return False + return commands.check(predicate) + + +def without_role(*role_ids: int) -> t.Callable: + """Check whether the invoking user does not have all of the roles specified in role_ids.""" + async def predicate(ctx: Context) -> bool: + if not ctx.guild: # Return False in a DM + log.debug( + f"{ctx.author} tried to use the '{ctx.command.name}' command from a DM. " + "This command is restricted by the without_role decorator. Rejecting request." + ) + return False + + author_roles = [role.id for role in ctx.author.roles] + check = all(role not in author_roles for role in role_ids) + log.debug( + f"{ctx.author} tried to call the '{ctx.command.name}' command. " + f"The result of the without_role check was {check}." + ) + return check + return commands.check(predicate) + + +def in_channel_check(*channels: int, bypass_roles: t.Container[int] = None) -> t.Callable[[Context], bool]: + """ + Checks that the message is in a whitelisted channel or optionally has a bypass role. + + If `in_channel_override` is present, check if it contains channels + and use them in place of the global whitelist. + """ + def predicate(ctx: Context) -> bool: + if not ctx.guild: + log.debug(f"{ctx.author} tried to use the '{ctx.command.name}' command from a DM.") + return True + if ctx.channel.id in channels: + log.debug( + f"{ctx.author} tried to call the '{ctx.command.name}' command " + f"and the command was used in a whitelisted channel." + ) + return True + + if bypass_roles and any(r.id in bypass_roles for r in ctx.author.roles): + log.debug( + f"{ctx.author} called the '{ctx.command.name}' command and " + f"had a role to bypass the in_channel check." + ) + return True + + if hasattr(ctx.command.callback, "in_channel_override"): + override = ctx.command.callback.in_channel_override + if override is None: + log.debug( + f"{ctx.author} called the '{ctx.command.name}' command " + f"and the command was whitelisted to bypass the in_channel check." + ) + return True + else: + if ctx.channel.id in override: + log.debug( + f"{ctx.author} tried to call the '{ctx.command.name}' command " + f"and the command was used in an overridden whitelisted channel." + ) + return True + + log.debug( + f"{ctx.author} tried to call the '{ctx.command.name}' command. " + f"The overridden in_channel check failed." + ) + channels_str = ', '.join(f"<#{c_id}>" for c_id in override) + raise InChannelCheckFailure( + f"Sorry, but you may only use this command within {channels_str}." + ) + + log.debug( + f"{ctx.author} tried to call the '{ctx.command.name}' command. " + f"The in_channel check failed." + ) + + channels_str = ', '.join(f"<#{c_id}>" for c_id in channels) + raise InChannelCheckFailure( + f"Sorry, but you may only use this command within {channels_str}." + ) + + return predicate + + +in_channel = commands.check(in_channel_check) + + +def override_in_channel(channels: t.Tuple[int] = None) -> t.Callable: + """ + Set command callback attribute for detection in `in_channel_check`. + + Override global whitelist if channels are specified. + + This decorator has to go before (below) below the `command` decorator. + """ + def inner(func: t.Callable) -> t.Callable: + func.in_channel_override = channels + return func + + return inner + + +def locked() -> t.Union[t.Callable, None]: + """ + Allows the user to only run one instance of the decorated command at a time. + + Subsequent calls to the command from the same author are ignored until the command has completed invocation. + + This decorator has to go before (below) the `command` decorator. + """ + def wrap(func: t.Callable) -> t.Union[t.Callable, None]: + func.__locks = WeakValueDictionary() + + @wraps(func) + async def inner(self: t.Callable, ctx: Context, *args, **kwargs) -> t.Union[t.Callable, None]: + lock = func.__locks.setdefault(ctx.author.id, Lock()) + if lock.locked(): + embed = Embed() + embed.colour = Colour.red() + + log.debug(f"User tried to invoke a locked command.") + embed.description = ( + "You're already using this command. Please wait until " + "it is done before you use it again." + ) + embed.title = random.choice(ERROR_REPLIES) + await ctx.send(embed=embed) + return + + async with func.__locks.setdefault(ctx.author.id, Lock()): + return await func(self, ctx, *args, **kwargs) + return inner + return wrap + + +def mock_in_debug(return_value: t.Any) -> t.Callable: + """ + Short-circuit function execution if in debug mode and return `return_value`. + + The original function name, and the incoming args and kwargs are DEBUG level logged + upon each call. This is useful for expensive operations, i.e. media asset uploads + that are prone to rate-limits but need to be tested extensively. + """ + def decorator(func: t.Callable) -> t.Callable: + @functools.wraps(func) + async def wrapped(*args, **kwargs) -> t.Any: + """Short-circuit and log if in debug mode.""" + if Client.debug: + log.debug(f"Function {func.__name__} called with args: {args}, kwargs: {kwargs}") + return return_value + return await func(*args, **kwargs) + return wrapped + return decorator -- cgit v1.2.3