[WelcomeCount] Implement most of #48

Avatars in welcome message is not yet a feature - will look at embed usage for this in the future.

Signed-off-by: Toby Harradine <[email protected]>

welcomecount/welcomecount.py

@@ -21,7 +21,7 @@
 # SOFTWARE.
 
 import datetime
-from typing import List
+from typing import List, Union
 
 import discord
 from redbot.core import Config, checks, commands
@@ -35,14 +35,15 @@
 )
 
 
-class WelcomeCount(getattr(commands, "Cog", object)):
+class WelcomeCount(commands.Cog):
     """A special welcome cog which keeps a daily count of new users.
 
     Idea came from Twentysix's version of Red on the official Red-DiscordBot
     server.
     """
 
     def __init__(self):
+        super().__init__()
         self.conf: Config = Config.get_conf(
             self, identifier=UNIQUE_ID, force_registration=True
         )
@@ -55,12 +56,12 @@ def __init__(self):
         self.conf.register_channel(
             enabled=False, last_message=None, welcome_msg=_DEFAULT_WELCOME
         )
-        self.conf.register_guild(count=0, day=None)
+        self.conf.register_guild(count=0, day=None, join_role=None)
 
-    @commands.group(invoke_without_command=True)
-    @commands.guild_only()
     @checks.admin_or_permissions(manage_guild=True)
-    async def wcount(self, ctx: commands.Context):
+    @commands.guild_only()
+    @commands.group(invoke_without_command=True, aliases=["wcount"])
+    async def welcomecount(self, ctx: commands.Context):
         """Manage settings for WelcomeCount."""
         if not ctx.invoked_subcommand:
             await ctx.send_help()
@@ -80,9 +81,9 @@ async def wcount(self, ctx: commands.Context):
             else:
                 await ctx.send(box("Disabled in this channel."))
 
-    @wcount.command(name="toggle", pass_context=True)
-    async def wcount_toggle(self, ctx: commands.Context):
-        """Enable/disable welcome messages in this channel."""
+    @welcomecount.command(name="toggle")
+    async def welcomecount_toggle(self, ctx: commands.Context):
+        """Toggle welcome messages in this channel."""
         channel: discord.TextChannel = ctx.channel
         settings = self.conf.channel(channel)
         now_enabled: bool = not await settings.enabled()
@@ -92,16 +93,17 @@ async def wcount_toggle(self, ctx: commands.Context):
             "".format("enabled" if now_enabled else "disabled")
         )
 
-    @wcount.command(name="message", pass_context=True, no_pm=True)
-    async def wcount_message(self, ctx: commands.Context, *, message: str):
+    @welcomecount.command(name="message")
+    async def welcomecount_message(self, ctx: commands.Context, *, message: str):
         """Set the bot's welcome message.
 
         This message can be formatted using these parameters:
             mention - Mention the user who joined
             username - The user's display name
             server - The name of the server
             count - The number of users who joined today.
-            plural - Empty if count is 1. 's' otherwise.
+            plural - Empty if `count` is 1. 's' otherwise.
+            total - The total number of users in the server.
         To format the welcome message with the above parameters, include them
         in your message surrounded by curly braces {}.
         """
@@ -116,13 +118,24 @@ async def wcount_message(self, ctx: commands.Context, *, message: str):
             "server": ctx.guild.name,
             "count": count,
             "plural": "" if count == 1 else "s",
+            "total": ctx.guild.member_count,
         }
-        await ctx.send("Welcome message set, sending a test message here...")
-        await ctx.send(message.format(**params))
+        try:
+            to_send = message.format(**params)
+        except KeyError as exc:
+            await ctx.send(
+                f"The welcome message cannot be formatted, because it contains an "
+                f"invalid placeholder `{{{exc.args[0]}}}`. See `{ctx.clean_prefix}help "
+                f"welcomecount message` for a list of valid placeholders."
+            )
+        else:
+            await ctx.send(
+                "Welcome message set, here's what it'll look like:\n\n" + to_send
+            )
 
-    @wcount.command(name="deletelast", pass_context=True)
-    async def wcount_deletelast(self, ctx: commands.Context):
-        """Enable/disable deleting previous welcome message in this channel.
+    @welcomecount.command(name="deletelast")
+    async def welcomecount_deletelast(self, ctx: commands.Context):
+        """Toggle deleting the previous welcome message in this channel.
 
         When enabled, the last message is deleted *only* if it was sent on
         the same day as the new welcome message.
@@ -136,11 +149,27 @@ async def wcount_deletelast(self, ctx: commands.Context):
             "".format("enabled" if now_deleting else "disabled")
         )
 
-    # Events
+    @welcomecount.command(name="joinrole")
+    async def welcomecount_joinrole(self, ctx: commands.Context, *, role: Union[discord.Role, str]):
+        """Set a role which a user must receive before they're welcomed.
 
-    @commands.Cog.listener()
-    async def on_member_join(self, member: discord.Member):
-        """Send the welcome message and update the last message."""
+        This means that, instead of the welcome message being sent when
+        the user joins the server, the welcome message will be sent when
+        they receive a particular role.
+
+        Use `[p]welcomecount joinrole disable` to revert to the default
+        behaviour.
+        """
+        if isinstance(role, discord.Role):
+            await self.conf.guild(ctx.guild).join_role.set(role.id)
+            await ctx.tick()
+        elif role.lower() == "disable":
+            await self.conf.guild(ctx.guild).join_role.clear()
+            await ctx.tick()
+        else:
+            await ctx.send(f'Role "{role}" not found.')
+
+    async def send_welcome_message(self, member: discord.Member) -> None:
         guild: discord.Guild = member.guild
         server_settings = self.conf.guild(guild)
         today: datetime.date = datetime.date.today()
@@ -182,7 +211,33 @@ async def on_member_join(self, member: discord.Member):
                 "server": guild.name,
                 "count": count,
                 "plural": "" if count == 1 else "s",
+                "total": guild.member_count,
             }
             welcome: str = await channel_settings.welcome_msg()
             msg: discord.Message = await channel.send(welcome.format(**params))
             await channel_settings.last_message.set(msg.id)
+
+    # Events
+
+    @commands.Cog.listener()
+    async def on_member_join(self, member: discord.Member):
+        """Send the welcome message and update the last message."""
+        if await self.conf.guild(member.guild).join_role() is None:
+            await self.send_welcome_message(member)
+
+    @commands.Cog.listener()
+    async def on_member_update(self, before: discord.Member, after: discord.Member):
+        join_role_id = await self.conf.guild(before.guild).join_role()
+        if join_role_id is None:
+            return
+
+        before_roles = frozenset(before.roles)
+        after_roles = frozenset(after.roles)
+        try:
+            added_role = next(iter(after_roles - before_roles))
+        except StopIteration:
+            # A role wasn't added
+            return
+
+        if added_role.id == join_role_id:
+            await self.send_welcome_message(after)