From ad504aba422cd5ca3eab783103d53f004227bbd1 Mon Sep 17 00:00:00 2001 From: Emily Doherty Date: Tue, 3 Dec 2024 16:59:50 -0800 Subject: [PATCH] Update the Commands with descriptions and implement remove-alert --- token_bot/tracker.py | 90 ++++++++++++++++++++++++++++++++++++++------ 1 file changed, 78 insertions(+), 12 deletions(-) diff --git a/token_bot/tracker.py b/token_bot/tracker.py index 73da91d..2cfd39d 100644 --- a/token_bot/tracker.py +++ b/token_bot/tracker.py @@ -4,7 +4,7 @@ from typing import Type, Dict, List import aiohttp from interactions import Extension, SlashContext, component_callback, \ - ComponentContext, StringSelectMenu, Message, Embed, EmbedField + ComponentContext, StringSelectMenu, Message, Embed, EmbedField, is_owner, check, StringSelectOption from interactions import Task, IntervalTrigger from interactions import slash_command, listen from interactions.api.events import Component @@ -100,7 +100,10 @@ class Tracker(Extension): self.update_data.start() - @slash_command() + @slash_command( + name="register", + description="Register with TokenBot for alerts on token price changes." + ) async def register(self, ctx: SlashContext): text = ("## Select a region to register with \n\n" "Please note: \n" @@ -109,16 +112,31 @@ class Tracker(Extension): "* You can remove all alerts and registration using ```/remove-registration```") await ctx.send(text, components=REGION_MENU, ephemeral=True) - @slash_command() + @slash_command( + name="remove-registration", + description="Remove all alerts and registration from TokenBot" + ) async def remove_registration(self, ctx: SlashContext): - await self._users.delete(ctx.user.id) + if await self._users.exists(ctx.user.id): + user = await self._users.get(ctx.user.id) + for alert in user.subscribed_alerts: + await self._alerts.remove_user(alert, user) + await self._users.delete(ctx.user.id) + await ctx.send("All alert subscriptions and user registration deleted", ephemeral=True) - @slash_command() + @slash_command( + name="exists", + description="Check if you are registered with TokenBot" + ) + @check(is_owner()) async def exists(self, ctx: SlashContext): await ctx.send(str(await self._users.exists(ctx.user.id)), ephemeral=True) - @slash_command() + @slash_command( + name="add-alert", + description="List all alerts you have signed up for" + ) async def add_alert(self, ctx: SlashContext): if not await self._users.exists(ctx.user.id): await ctx.send("You are not registered with any region\n" @@ -145,18 +163,33 @@ class Tracker(Extension): alert = Alert(alert_type, flavor, user.region) if not await self._users.is_subscribed(user, alert): await self._users.add_alert(user, alert) - await self._alerts.add_user(user, alert) + await self._alerts.add_user(alert, user) await ctx.send("Successfully added alert", ephemeral=True) else: await ctx.send("You are already subscribed to this alert", ephemeral=True) - @slash_command() + @slash_command( + name="remove-alert", + description="Remove an alert you have signed up for" + ) async def remove_alert(self, ctx: SlashContext): - await ctx.send("This is not implemented yet, use /remove-registration for the time being", ephemeral=True) + user = await self._users.get(ctx.user.id) + try: + alert = await self.remove_alert_select_menu(ctx, user) + except TimeoutError: + return + else: + await self._users.remove_alert(user, alert) + await self._alerts.remove_user(alert, user) + await ctx.send("Successfully removed alert", ephemeral=True) - @slash_command() + + @slash_command( + name="list-alerts", + description="List all alerts you have signed up for" + ) async def list_alerts(self, ctx: SlashContext): alerts = await self._users.list_alerts(ctx.user.id) alerts_str = f"You have {len(alerts)} out of 25 maximum alerts registered" @@ -174,8 +207,6 @@ class Tracker(Extension): # Callbacks Commands # ################################### - # TODO: export to a separate file if possible - @component_callback('flavor_menu') async def flavor_menu(self, ctx: ComponentContext): await ctx.send(f"Selected Flavor: {ctx.values[0]}", ephemeral=True) @@ -188,6 +219,10 @@ class Tracker(Extension): async def alert_menu(self, ctx: ComponentContext): await ctx.send(f"Selected Alert: {ctx.values[0]}", ephemeral=True) + @component_callback('remove_alert_menu') + async def remove_alert_menu(self, ctx: ComponentContext): + await ctx.send(f"You have selected to remove the following alert: {ctx.values[0].title()}", ephemeral=True) + @component_callback('region_menu') async def region_menu(self, ctx: ComponentContext): user = User(ctx.user.id, Region(ctx.values[0].lower()), subscribed_alerts=[]) @@ -210,6 +245,37 @@ class Tracker(Extension): # Helper Functions # ################################### + async def remove_alert_select_menu(self, ctx: SlashContext, user: User): + alerts_by_flavor = await gather_alerts_by_flavor(user.subscribed_alerts) + select_options: List[StringSelectOption] = [] + for flavor in alerts_by_flavor: + for alert in alerts_by_flavor[flavor]: + select_options.append(StringSelectOption( + label=f"{alert.flavor.name.lower().title()} {alert.to_human_string()}", + value=f"{alert.flavor.name.lower()} {alert.to_human_string()}" + )) + menu = StringSelectMenu( + select_options, + placeholder="Select an alert to remove", + custom_id="remove_alert_menu" + ) + message = await ctx.send("Select an alert to remove", components=menu, ephemeral=True) + + try: + alert_component: Component = await self.bot.wait_for_component(messages=message, + components=menu, timeout=30) + except TimeoutError: + menu.disabled = True + await message.edit(context=ctx, components=menu, content="Timed out") + raise TimeoutError + else: + menu.disabled = True + await message.edit(context=ctx, components=menu) + selection_split = alert_component.ctx.values[0].split(" ") + flavor = Flavor[selection_split[0].upper()] + alert_type = AlertType.from_str(' '.join(selection_split[1:])) + return Alert(alert_type, flavor, user.region) + async def flavor_select_menu(self, ctx: SlashContext) -> Type[Flavor]: flavor_menu = copy.deepcopy(FLAVOR_MENU)