diff --git a/timerole/timerole.py b/timerole/timerole.py index b823dd3..f50b7af 100644 --- a/timerole/timerole.py +++ b/timerole/timerole.py @@ -26,7 +26,11 @@ class Timerole(Cog): @checks.guildowner() @commands.guild_only() async def runtimerole(self, ctx: commands.Context): - """Trigger the daily timerole""" + """ + Trigger the daily timerole + + Useful for troubleshooting the initial setup + """ await self.timerole_update() await ctx.send("Success") @@ -46,12 +50,30 @@ class Timerole(Cog): """Add a role to be added after specified time on server""" guild = ctx.guild - to_set = {"days": days} + to_set = {"days": days, "remove": False} if requiredroles: to_set["required"] = [r.id for r in requiredroles] await self.config.guild(guild).roles.set_raw(role.id, value=to_set) - await ctx.send("Time Role for {0} set to {1} days".format(role.name, days)) + await ctx.maybe_send_embed("Time Role for {0} set to {1} days until added".format(role.name, days)) + + @timerole.command() + async def removerole( + self, ctx: commands.Context, role: discord.Role, days: int, *requiredroles: discord.Role + ): + """ + Add a role to be removed after specified time on server + + Useful with an autorole cog + """ + guild = ctx.guild + + to_set = {"days": days, "remove": True} + if requiredroles: + to_set["required"] = [r.id for r in requiredroles] + + await self.config.guild(guild).roles.set_raw(role.id, value=to_set) + await ctx.maybe_send_embed("Time Role for {0} set to {1} days until removed".format(role.name, days)) @timerole.command() async def channel(self, ctx: commands.Context, channel: discord.TextChannel): @@ -62,8 +84,8 @@ class Timerole(Cog): await ctx.send("Announce channel set to {0}".format(channel.mention)) @timerole.command() - async def removerole(self, ctx: commands.Context, role: discord.Role): - """Removes a role from being added after specified time""" + async def delrole(self, ctx: commands.Context, role: discord.Role): + """Deletes a role from being added/removed after specified time""" guild = ctx.guild await self.config.guild(guild).roles.set_raw(role.id, value=None) @@ -93,6 +115,7 @@ class Timerole(Cog): async def timerole_update(self): for guild in self.bot.guilds: addlist = [] + removelist = [] role_dict = await self.config.guild(guild).roles() if not any(role_data for role_data in role_dict.values()): # No roles @@ -101,39 +124,56 @@ class Timerole(Cog): for member in guild.members: has_roles = [r.id for r in member.roles] - get_roles = [int(rID) for rID, r_data in role_dict.items() if r_data is not None] - - check_roles = set(get_roles) - set(has_roles) + add_roles = [int(rID) for rID, r_data in role_dict.items() if r_data is not None and not r_data["remove"]] + remove_roles = [int(rID) for rID, r_data in role_dict.items() if r_data is not None and r_data["remove"]] - for role_id in check_roles: - # Check for required role - if "required" in role_dict[str(role_id)]: - if not set(role_dict[str(role_id)]["required"]) & set(has_roles): - # Doesn't have required role - continue + check_add_roles = set(add_roles) - set(has_roles) + check_remove_roles = set(remove_roles) & set(has_roles) - if ( - member.joined_at + timedelta(days=role_dict[str(role_id)]["days"]) - <= datetime.today() - ): - # Qualifies - addlist.append((member, role_id)) + await self.check_required_and_date(addlist, check_add_roles, has_roles, member, role_dict) + await self.check_required_and_date(removelist, check_remove_roles, has_roles, member, role_dict) channel = await self.config.guild(guild).announce() if channel is not None: channel = guild.get_channel(channel) title = "**These members have received the following roles**\n" - results = "" - for member, role_id in addlist: - role = discord.utils.get(guild.roles, id=role_id) - await member.add_roles(role, reason="Timerole") + await self.announce_roles(title, addlist, channel, guild, to_add=True) + title = "**These members have lost the following roles**\n" + await self.announce_roles(title, removelist, channel, guild, to_add=False) + + async def announce_roles(self, title, role_list, channel, guild, to_add: True): + results = "" + for member, role_id in role_list: + role = discord.utils.get(guild.roles, id=role_id) + try: + if to_add: + await member.add_roles(role, reason="Timerole") + else: + await member.remove_roles(role, reason="Timerole") + except discord.Forbidden: + results += "{} : {} **(Failed)**\n".format(member.display_name, role.name) + else: results += "{} : {}\n".format(member.display_name, role.name) - - if channel is not None and results: - await channel.send(title) - for page in pagify(results, shorten_by=50): - await channel.send(page) + if channel is not None and results: + await channel.send(title) + for page in pagify(results, shorten_by=50): + await channel.send(page) + + async def check_required_and_date(self, role_list, check_roles, has_roles, member, role_dict): + for role_id in check_roles: + # Check for required role + if "required" in role_dict[str(role_id)]: + if not set(role_dict[str(role_id)]["required"]) & set(has_roles): + # Doesn't have required role + continue + + if ( + member.joined_at + timedelta(days=role_dict[str(role_id)]["days"]) + <= datetime.today() + ): + # Qualifies + role_list.append((member, role_id)) async def check_day(self): while self is self.bot.get_cog("Timerole"):