[Streams] Save/Remove/Check custom messages for each streamer.

redbot/cogs/streams/streams.py

@@ -77,6 +77,7 @@ def __init__(self, bot: Red):
         self.bot: Red = bot
 
         self.streams: List[Stream] = []
+        self.streamer_info = {}
         self.task: Optional[asyncio.Task] = None
 
         self.yt_cid_pattern = re.compile("^UC[-_A-Za-z0-9]{21}[AQgw]$")
@@ -97,6 +98,7 @@ async def cog_load(self) -> None:
             await self.move_api_keys()
             await self.get_twitch_bearer_token()
             self.streams = await self.load_streams()
+            self.streamer_info = await self.load_streamer_info()
             self.task = asyncio.create_task(self._stream_alerts())
         except Exception as error:
             log.exception("Failed to initialize Streams cog:", exc_info=error)
@@ -508,7 +510,7 @@ async def _streamset_refresh_timer(self, ctx: commands.Context, refresh_time: in
 
         await self.config.refresh_timer.set(refresh_time)
         await ctx.send(
-            _("Refresh timer set to {refresh_time} seconds".format(refresh_time=refresh_time))
+            _("Refresh timer set to {refresh_time} seconds").format(refresh_time=refresh_time)
         )
 
     @streamset.command()
@@ -600,13 +602,176 @@ async def without_mention(self, ctx: commands.Context, *, message: str):
         await self.config.guild(guild).live_message_nomention.set(message)
         await ctx.send(_("Stream alert message set!"))
 
+    @message.command(name="check")
+    @commands.guild_only()
+    async def check_messages(self, ctx: commands.Context, streamer_name: str):
+        """
+        Use the command `[p]streamset message check <streamer_name>`
+
+        Red will respond with the message(s) associated with that streamer.
+        """
+
+        for stream in self.streams:
+            if stream.name.lower() == streamer_name.lower():
+                try:
+                    await ctx.send(
+                        _("`{}`'s nomention message is set as: `{}`").format(
+                            streamer_name, stream.nomention_message
+                        )
+                    )
+                except AttributeError:
+                    await ctx.send(
+                        _("`{}` does not have a nomention message.").format(streamer_name)
+                    )
+                try:
+                    await ctx.send(
+                        _("`{}`'s mention message is set as: `{}` to `{}`").format(
+                            streamer_name, stream.mention_message, stream.who_to_mention
+                        )
+                    )
+                except AttributeError:
+                    await ctx.send(
+                        _("`{}` does not have a mention message.").format(streamer_name)
+                    )
+                return
+
+        await ctx.send(
+            _(
+                "Streamer `{}` is not registered and therefore does not have any custom messages."
+            ).format(streamer_name)
+        )
+        return
+
+    # @checks.mod_or_permissions(manage_channels=True)
+    @message.command(name="streamer")
+    @commands.guild_only()
+    async def custom_message(
+        self,
+        ctx: commands.Context,
+        streamer_name: str,
+        mention: str,
+        to_mention: str = None,
+        *,
+        msg: str,
+    ):
+        """Set custom stream alert message for already-registered streamer.
+
+        Use the command `[p]streamset message streamer <streamer_name> <mention|nomention> [if_mention: audience] <message>`
+
+        Streamer must be already registered.
+        Command can be run by moderator.
+        Can only be used in server.
+        """
+
+        # Change in the print statements way below too
+        streams_list = defaultdict(list)
+
+        # Verifying that the user is setting a message for an existing streamer
+        not_found = True
+        stream_id = ""
+        for stream in self.streams:
+            if stream.name.lower() == streamer_name.lower():
+                stream_id = stream.id
+                not_found = False
+                if mention == "mention":
+                    stream.mention_message = msg
+                    stream.who_to_mention = "@{}".format(to_mention)
+                    await ctx.send(
+                        _(
+                            "Custom message for streamer `{streamer}` set to mention @\u200b{mention}."
+                        ).format(streamer=streamer_name, mention=to_mention)
+                    )
+                    await self.save_streams()
+                    break
+                elif mention == "nomention":
+                    msg = to_mention + " " + msg
+                    stream.nomention_message = msg
+                    await ctx.send(
+                        _("Custom nomention message streamer `{streamer}` set").format(
+                            streamer=streamer_name
+                        )
+                    )
+                    await self.save_streams()
+                    break
+                else:
+                    await ctx.send_help()
+                    break
+
+        if not_found:
+            await ctx.send(
+                _(
+                    "Streamer `{streamer}` not registered, please look at `[p]streamalert help`"
+                ).format(streamer=streamer_name)
+            )
+            return
+
+    @message.command(name="remove")
+    @commands.guild_only()
+    async def remove_message(self, ctx: commands.Context, name: str, mention: str):
+        """Remove custom stream alert message for already-registered streamer.
+
+        Use the command `[p]streamset message remove <streamer> <mention|nomention|both>`
+
+        Streamer must be already registered.
+        Command can be run by moderator.
+        Can only be used in server.
+
+        Note: this reverts the message for one streamer only, not all.
+        """
+
+        streams_list = defaultdict(list)
+        guild_channels_ids = [c.id for c in ctx.guild.channels]
+
+        not_found = True
+        for stream in self.streams:
+            if stream.name.lower() == name.lower():
+                not_found = False
+                if mention == "mention":
+                    if not hasattr(stream, "mention_message"):
+                        break
+                    del stream.mention_message
+                    del stream.who_to_mention
+                    break
+                elif mention == "nomention":
+                    if not hasattr(stream, "nomention_message"):
+                        break
+                    del stream.nomention_message
+                    break
+                elif mention == "both":
+                    if hasattr(stream, "nomention_message"):
+                        del stream.nomention_message
+                    if hasattr(stream, "mention_message"):
+                        del stream.mention_message
+                        del stream.who_to_mention
+                    break
+                else:
+                    await ctx.send_help()
+
+        if not_found:
+            await ctx.send(
+                _("Streamer `{}` not registered. No message(s) to remove.").format(name)
+            )
+            return
+
+        await self.save_streams()
+        await ctx.send(_("Removed message(s) from streamer `{}`").format(name))
+
     @message.command(name="clear")
     @commands.guild_only()
     async def clear_message(self, ctx: commands.Context):
         """Reset the stream alert messages in this server."""
         guild = ctx.guild
         await self.config.guild(guild).live_message_mention.set(False)
         await self.config.guild(guild).live_message_nomention.set(False)
+
+        for stream in self.streams:
+            if hasattr(stream, "nomention_message"):
+                del stream.nomention_message
+            if hasattr(stream, "mention_message"):
+                del stream.mention_message
+                del stream.who_to_mention
+
+        await self.save_streams()
         await ctx.send(_("Stream alerts in this server will now use the default alert message."))
 
     @streamset.group()
@@ -899,9 +1064,16 @@ async def check_streams(self):
                         )
 
                         if mention_str:
-                            if guild_data["live_message_mention"]:
-                                # Stop bad things from happening here...
-                                content = guild_data["live_message_mention"]
+                            alert_msg = await self.config.guild(
+                                channel.guild
+                            ).live_message_mention()
+                            if hasattr(stream, "mention_message"):
+                                content = _("{mention}, {mention_message}").format(
+                                    mention=stream.who_to_mention,
+                                    mention_message=stream.mention_message,
+                                )
+                            elif alert_msg:
+                                content = alert_msg  # Stop bad things from happening here...
                                 content = content.replace(
                                     "{stream.name}", str(stream.name)
                                 )  # Backwards compatibility
@@ -920,9 +1092,11 @@ async def check_streams(self):
                                     ),
                                 )
                         else:
-                            if guild_data["live_message_nomention"]:
-                                # Stop bad things from happening here...
-                                content = guild_data["live_message_nomention"]
+                            alert_msg = await self.config.guild(
+                                channel.guild
+                            ).live_message_nomention()
+                            if alert_msg:
+                                content = alert_msg  # Stop bad things from happening here...
                                 content = content.replace(
                                     "{stream.name}", str(stream.name)
                                 )  # Backwards compatibility
@@ -1003,6 +1177,29 @@ async def load_streams(self):
             _class = getattr(_streamtypes, raw_stream["type"], None)
             if not _class:
                 continue
+            if "nomention_message" in raw_stream:
+                if raw_stream["id"] not in self.streamer_info:
+                    self.streamer_info[raw_stream["id"]] = {
+                        "nomention_message": raw_stream["nomention_message"]
+                    }
+                else:
+                    self.streamer_info[raw_stream["id"]]["nomention_message"] = raw_stream[
+                        "nomention_message"
+                    ]
+            if "mention_message" in raw_stream:
+                if raw_stream["id"] not in self.streamer_info:
+                    self.streamer_info[raw_stream["id"]] = {
+                        "mention_message": raw_stream["mention_message"],
+                        "who_to_mention": raw_stream["who_to_mention"],
+                    }
+                else:
+                    self.streamer_info[raw_stream["id"]]["mention_message"] = raw_stream[
+                        "mention_message"
+                    ]
+                    self.streamer_info[raw_stream["id"]]["who_to_mention"] = raw_stream[
+                        "who_to_mention"
+                    ]
+
             token = await self.bot.get_shared_api_tokens(_class.token_name)
             if token:
                 if _class.__name__ == "TwitchStream":
@@ -1017,6 +1214,17 @@ async def load_streams(self):
 
         return streams
 
+    async def load_streamer_info(self):
+        for current_stream in self.streams:
+            id = current_stream.id
+            if id in self.streamer_info:
+                info = self.streamer_info[id]
+                if "nomention_message" in info:
+                    current_stream.nomention_message = info["nomention_message"]
+                if "mention_message" in info:
+                    current_stream.mention_message = info["mention_message"]
+                    current_stream.who_to_mention = info["who_to_mention"]
+
     async def save_streams(self):
         raw_streams = []
         for stream in self.streams: