diff --git a/token_bot/tracker.py b/token_bot/tracker.py index 204d32a..67a7e82 100644 --- a/token_bot/tracker.py +++ b/token_bot/tracker.py @@ -85,13 +85,13 @@ class Tracker(Extension): description=f"Hello, you requested to be sent an alert when the price of the World of Warcraft " f"token reaches a certain value.\n\n" f"As a reminder, you can remove an alert via ```/remove-alert```\n" - f"or you can remove all registrations via ```/remove-registration```\n\n", + f"or you can remove all alerts and user data via ```/remove-registration```\n\n", ) ] alerts_by_flavor = await gather_alerts_by_flavor(users_alerts[user]) for flavor in alerts_by_flavor: embeds.append( - await self.render_alert_flavor(alerts_by_flavor[flavor], user=user) + await self._render_alert_flavor(alerts_by_flavor[flavor], user=user) ) await discord_user.send(embeds=embeds) @@ -118,7 +118,7 @@ class Tracker(Extension): @slash_command( name="register", - description="Register with TokenBot for alerts on token price changes.", + description="Register with a new GoblinBot Region for alerts on token price changes.", ) async def register(self, ctx: SlashContext): text = ( @@ -126,14 +126,14 @@ class Tracker(Extension): "Please note: \n" "* You can only be registered with one region at a time \n" "* Changing your region will remove all previous alerts you have signed up for \n" - "* You can remove all alerts and registration using ```/remove-registration```" + "* You can remove all alerts and user data using ```/remove-registration```" ) menu = copy.deepcopy(REGION_MENU) await ctx.send(text, components=menu, ephemeral=True) @slash_command( name="remove-registration", - description="Remove all alerts and registration from TokenBot", + description="Remove all alerts and registration from GoblinBot", ) async def remove_registration(self, ctx: SlashContext): if await self._users.exists(ctx.user.id): @@ -142,12 +142,10 @@ class Tracker(Extension): 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 - ) + await ctx.send("All alert subscriptions and user data deleted", ephemeral=True) @slash_command( - name="exists", description="Check if you are registered with TokenBot" + name="exists", description="Check if you are registered with GoblinBot" "" ) @check(is_owner()) async def exists(self, ctx: SlashContext): @@ -166,12 +164,11 @@ class Tracker(Extension): @slash_command(name="add-alert", description="Add an alert listener") 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" - "Please register with /register before adding alerts", - ephemeral=True, - ) - return + try: + await self.region_select_menu(ctx) + except TimeoutError: + return + user = await self._users.get(ctx.user.id) try: @@ -206,12 +203,7 @@ class Tracker(Extension): name="remove-alert", description="Remove an alert you have signed up for" ) async def remove_alert(self, ctx: SlashContext): - if not await self._users.exists(ctx.user.id): - await ctx.send( - "You are not registered with any region\n" - "Please register with /register before adding alerts", - ephemeral=True, - ) + if not self._user_is_registered(ctx): return user = await self._users.get(ctx.user.id) alerts = await self._users.list_alerts(user) @@ -234,12 +226,7 @@ class Tracker(Extension): name="list-alerts", description="List all alerts you have signed up for" ) async def list_alerts(self, ctx: SlashContext): - if not await self._users.exists(ctx.user.id): - await ctx.send( - "You are not registered with any region\n" - "Please register with /register before adding alerts", - ephemeral=True, - ) + if not await self._user_is_registered(ctx): return user = await self._users.get(ctx.user.id) alerts = await self._users.list_alerts(user) @@ -249,7 +236,7 @@ class Tracker(Extension): alerts_str = f"You have {len(alerts)} out of 25 maximum alerts registered" embeds = [ Embed( - title="List of TokenBot Tracker Alerts", + title="List of GoblinBot Tracker Alerts", color=0x0000B1, description=alerts_str, ) @@ -257,7 +244,7 @@ class Tracker(Extension): alerts_by_flavor = await gather_alerts_by_flavor(alerts) for flavor in alerts_by_flavor: embeds.append( - await self.render_alert_flavor(alerts_by_flavor[flavor], user=user) + await self._render_alert_flavor(alerts_by_flavor[flavor], user=user) ) await ctx.send(embeds=embeds, ephemeral=True) @@ -285,18 +272,14 @@ class Tracker(Extension): ) @component_callback("region_menu") - async def region_menu(self, ctx: ComponentContext): - user = User(ctx.user.id, Region(ctx.values[0].lower()), subscribed_alerts=[]) - await self._users.add(user) - discord_user = await self.bot.fetch_user(user.user_id) + async def region_menu_cb(self, ctx: ComponentContext): + discord_user = await self.bot.fetch_user(ctx.user.id) await discord_user.send( - "You have successfully registered with TokenBot!\n" - "Most interactions will happen in direct messages with TokenBot here.\n" - "You can remove your registration and alerts at any time using ```/remove-registration```\n" - ) - await ctx.send( - f"Successfully registered with the {ctx.values[0]} region", ephemeral=True + "You have successfully registered your region with GoblinBot!\n" + "Most interactions will happen in direct messages with GoblinBot here.\n" + "You can remove your user data and alerts at any time using ```/remove-registration```\n" ) + await ctx.defer(edit_origin=True) @component_callback("high_alert_button") async def high_alert_button(self, ctx: ComponentContext): @@ -364,6 +347,41 @@ class Tracker(Extension): alert_type = AlertType.from_str(" ".join(selection_split[1:])) return Alert(alert_type, flavor, user.region) + async def region_select_menu(self, ctx: SlashContext, user: User | None = None): + region_menu = copy.deepcopy(REGION_MENU) + region_menu_str = str() + + if user is None: + region_menu_str += "You are not currently registered with a region, please select a region to register with.\n" + region_menu_str += ( + "* You can only be registered with one region at a time.\n" + "* Registering for a new region will remove your old region's registration.\n" + ) + region_message = await ctx.send( + region_menu_str, + components=region_menu, + ephemeral=True, + ) + try: + region_component = await self.bot.wait_for_component( + messages=region_message, components=region_menu, timeout=30 + ) + except TimeoutError: + region_menu.disabled = True + await region_message.edit( + context=ctx, components=region_menu, content="Timed out" + ) + raise TimeoutError + else: + region_menu.disabled = True + region = Region(region_component.ctx.values[0].lower()) + user = User(ctx.user.id, region, subscribed_alerts=[]) + await asyncio.gather( + self._users.add(user), + region_message.edit(context=ctx, components=region_menu), + ) + return region + async def flavor_select_menu(self, ctx: SlashContext) -> Type[Flavor]: flavor_menu = copy.deepcopy(FLAVOR_MENU) @@ -422,6 +440,7 @@ class Tracker(Extension): raise TimeoutError else: menu.disabled = True + await component.ctx.defer(edit_origin=True) await message.edit(context=ctx, components=menu) return AlertType.from_str(component.ctx.values[0]) @@ -443,7 +462,17 @@ class Tracker(Extension): ) return await self._alert_select_menu_handler(ctx, low_menu, low_message) - async def render_alert_flavor( + async def _user_is_registered(self, ctx: SlashContext) -> bool: + if not await self._users.exists(ctx.user.id): + await ctx.send( + "You are not registered with any region\n" + "Please add an alert to get started ```/add-alert```", + ephemeral=True, + ) + return False + return True + + async def _render_alert_flavor( self, alerts: List[Alert], user: User | None = None ) -> Embed: region = alerts[0].region