From 0010d2d332f64bebd0ca4eadeac682c7f83014f9 Mon Sep 17 00:00:00 2001 From: MarkKoz Date: Sat, 18 Jul 2020 13:48:56 -0700 Subject: Decorators: wrap mutually_exclusive to use get_arg_value Instead of taking a callable, this wrapper just takes a name or position to get the resource ID. --- bot/cogs/reminders.py | 10 ++++------ bot/decorators.py | 17 ++++++++++++++++- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/bot/cogs/reminders.py b/bot/cogs/reminders.py index 1a0a9d303..30f7c8876 100644 --- a/bot/cogs/reminders.py +++ b/bot/cogs/reminders.py @@ -4,7 +4,6 @@ import random import textwrap import typing as t from datetime import datetime, timedelta -from functools import partial from operator import itemgetter import discord @@ -15,10 +14,9 @@ from discord.ext.commands import Cog, Context, Greedy, group from bot.bot import Bot from bot.constants import Guild, Icons, MODERATION_ROLES, POSITIVE_REPLIES, STAFF_ROLES from bot.converters import Duration -from bot.decorators import mutually_exclusive +from bot.decorators import mutually_exclusive_arg from bot.pagination import LinePaginator from bot.utils.checks import without_role_check -from bot.utils.function import get_arg_value from bot.utils.messages import send_denial from bot.utils.scheduling import Scheduler from bot.utils.time import humanize_delta @@ -168,7 +166,7 @@ class Reminders(Cog): log.trace(f"Scheduling new task #{reminder['id']}") self.schedule_reminder(reminder) - @mutually_exclusive(NAMESPACE, lambda args: get_arg_value("reminder", args)["id"]) + @mutually_exclusive_arg(NAMESPACE, "reminder", itemgetter("id")) async def send_reminder(self, reminder: dict, late: relativedelta = None) -> None: """Send the reminder.""" is_valid, user, channel = self.ensure_valid_reminder(reminder) @@ -375,7 +373,7 @@ class Reminders(Cog): mention_ids = [mention.id for mention in mentions] await self.edit_reminder(ctx, id_, {"mentions": mention_ids}) - @mutually_exclusive(NAMESPACE, partial(get_arg_value, "id_")) + @mutually_exclusive_arg(NAMESPACE, "id_") async def edit_reminder(self, ctx: Context, id_: int, payload: dict) -> None: """Edits a reminder with the given payload, then sends a confirmation message.""" reminder = await self._edit_reminder(id_, payload) @@ -393,7 +391,7 @@ class Reminders(Cog): await self._reschedule_reminder(reminder) @remind_group.command("delete", aliases=("remove", "cancel")) - @mutually_exclusive(NAMESPACE, partial(get_arg_value, "id_")) + @mutually_exclusive_arg(NAMESPACE, "id_") async def delete_reminder(self, ctx: Context, id_: int) -> None: """Delete one of your active reminders.""" await self.bot.api_client.delete(f"bot/reminders/{id_}") diff --git a/bot/decorators.py b/bot/decorators.py index cffd97440..c9e4a0560 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -5,7 +5,7 @@ import random import typing as t from collections import defaultdict from contextlib import suppress -from functools import wraps +from functools import partial, wraps from weakref import WeakValueDictionary from discord import Colour, Embed, Member, NotFound @@ -150,6 +150,21 @@ def mutually_exclusive(namespace: t.Hashable, resource_id: ResourceId) -> t.Call return decorator +def mutually_exclusive_arg( + namespace: t.Hashable, + name_or_pos: function.Argument, + func: t.Callable[[t.Any], _IdCallableReturn] = None +) -> t.Callable: + """ + Apply `mutually_exclusive` using the value of the arg at the given name/position as the ID. + + `func` is an optional callable or awaitable which will return the ID given the argument value. + See `mutually_exclusive` docs for more information. + """ + decorator_func = partial(mutually_exclusive, namespace) + return function.get_arg_value_wrapper(decorator_func, name_or_pos, func) + + def redirect_output(destination_channel: int, bypass_roles: t.Container[int] = None) -> t.Callable: """ Changes the channel in the context of the command to redirect the output to a certain channel. -- cgit v1.2.3