Merge pull request #151 from bobloy/timerole_develop

Timerole develop
pull/152/head
bobloy 4 years ago committed by GitHub
commit 60806fb19c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -1,6 +1,7 @@
import asyncio import asyncio
import logging import logging
from datetime import datetime, timedelta from datetime import datetime, timedelta
from typing import Optional
import discord import discord
from redbot.core import Config, checks, commands from redbot.core import Config, checks, commands
@ -19,6 +20,15 @@ async def sleep_till_next_hour():
await asyncio.sleep((next_hour - datetime.utcnow()).seconds) await asyncio.sleep((next_hour - datetime.utcnow()).seconds)
async def announce_to_channel(channel, results, title):
if channel is not None and results:
await channel.send(title)
for page in pagify(results, shorten_by=50):
await channel.send(page)
elif results: # Channel is None, log the results
log.info(results)
class Timerole(Cog): class Timerole(Cog):
"""Add roles to users based on time on server""" """Add roles to users based on time on server"""
@ -27,10 +37,15 @@ class Timerole(Cog):
self.bot = bot self.bot = bot
self.config = Config.get_conf(self, identifier=9811198108111121, force_registration=True) self.config = Config.get_conf(self, identifier=9811198108111121, force_registration=True)
default_global = {} default_global = {}
default_guild = {"announce": None, "roles": {}} default_guild = {"announce": None, "reapply": True, "roles": {}}
default_rolemember = {"had_role": False, "check_again_time": None}
self.config.register_global(**default_global) self.config.register_global(**default_global)
self.config.register_guild(**default_guild) self.config.register_guild(**default_guild)
self.config.init_custom("RoleMember", 2)
self.config.register_custom("RoleMember", **default_rolemember)
self.updating = asyncio.create_task(self.check_hour()) self.updating = asyncio.create_task(self.check_hour())
async def red_delete_data_for_user(self, **kwargs): async def red_delete_data_for_user(self, **kwargs):
@ -49,11 +64,14 @@ class Timerole(Cog):
Useful for troubleshooting the initial setup Useful for troubleshooting the initial setup
""" """
async with ctx.typing(): async with ctx.typing():
pre_run = datetime.utcnow()
await self.timerole_update() await self.timerole_update()
after_run = datetime.utcnow()
await ctx.tick() await ctx.tick()
await ctx.maybe_send_embed(f"Took {after_run-pre_run} seconds")
@commands.group() @commands.group()
@checks.mod_or_permissions(administrator=True) @checks.mod_or_permissions(administrator=True)
@commands.guild_only() @commands.guild_only()
@ -84,9 +102,7 @@ class Timerole(Cog):
await self.config.guild(guild).roles.set_raw(role.id, value=to_set) await self.config.guild(guild).roles.set_raw(role.id, value=to_set)
await ctx.maybe_send_embed( await ctx.maybe_send_embed(
"Time Role for {0} set to {1} days and {2} hours until added".format( f"Time Role for {role.name} set to {days} days and {hours} hours until added"
role.name, days, hours
)
) )
@timerole.command() @timerole.command()
@ -114,18 +130,27 @@ class Timerole(Cog):
await self.config.guild(guild).roles.set_raw(role.id, value=to_set) await self.config.guild(guild).roles.set_raw(role.id, value=to_set)
await ctx.maybe_send_embed( await ctx.maybe_send_embed(
"Time Role for {0} set to {1} days and {2} hours until removed".format( f"Time Role for {role.name} set to {days} days and {hours} hours until removed"
role.name, days, hours
)
) )
@timerole.command() @timerole.command()
async def channel(self, ctx: commands.Context, channel: discord.TextChannel): async def channel(self, ctx: commands.Context, channel: Optional[discord.TextChannel] = None):
"""Sets the announce channel for role adds""" """Sets the announce channel for role adds"""
guild = ctx.guild guild = ctx.guild
if channel is None:
await self.config.guild(guild).announce.clear()
await ctx.maybe_send_embed(f"Announce channel has been cleared")
else:
await self.config.guild(guild).announce.set(channel.id) await self.config.guild(guild).announce.set(channel.id)
await ctx.send("Announce channel set to {0}".format(channel.mention)) await ctx.send(f"Announce channel set to {channel.mention}")
@timerole.command()
async def reapply(self, ctx: commands.Context):
"""Toggle reapplying roles if the member loses it somehow. Defaults to True"""
guild = ctx.guild
current_setting = await self.config.guild(guild).reapply()
await self.config.guild(guild).reapply.set(not current_setting)
await ctx.maybe_send_embed(f"Reapplying roles is now set to: {not current_setting}")
@timerole.command() @timerole.command()
async def delrole(self, ctx: commands.Context, role: discord.Role): async def delrole(self, ctx: commands.Context, role: discord.Role):
@ -133,7 +158,8 @@ class Timerole(Cog):
guild = ctx.guild guild = ctx.guild
await self.config.guild(guild).roles.set_raw(role.id, value=None) await self.config.guild(guild).roles.set_raw(role.id, value=None)
await ctx.send("{0} will no longer be applied".format(role.name)) await self.config.custom("RoleMember", role.id).clear()
await ctx.maybe_send_embed(f"{role.name} will no longer be applied")
@timerole.command() @timerole.command()
async def list(self, ctx: commands.Context): async def list(self, ctx: commands.Context):
@ -153,89 +179,197 @@ class Timerole(Cog):
str(discord.utils.get(guild.roles, id=int(new_id))) str(discord.utils.get(guild.roles, id=int(new_id)))
for new_id in r_data["required"] for new_id in r_data["required"]
] ]
out += "{} | {} days | requires: {}\n".format(str(role), r_data["days"], r_roles) out += f"{role} | {r_data['days']} days | requires: {r_roles}\n"
await ctx.maybe_send_embed(out) await ctx.maybe_send_embed(out)
async def timerole_update(self): async def timerole_update(self):
async for guild in AsyncIter(self.bot.guilds): utcnow = datetime.utcnow()
all_guilds = await self.config.all_guilds()
# all_mrs = await self.config.custom("RoleMember").all()
# log.debug(f"Begin timerole update")
for guild in self.bot.guilds:
guild_id = guild.id
if guild_id not in all_guilds:
log.debug(f"Guild has no configured settings: {guild}")
continue
add_results = ""
remove_results = ""
reapply = all_guilds[guild_id]["reapply"]
role_dict = all_guilds[guild_id]["roles"]
if not any(role_data for role_data in role_dict.values()): # No roles
log.debug(f"No roles are configured for guild: {guild}")
continue
# all_mr = await self.config.all_custom("RoleMember")
# log.debug(f"{all_mr=}")
async for member in AsyncIter(guild.members, steps=10):
addlist = [] addlist = []
removelist = [] removelist = []
role_dict = await self.config.guild(guild).roles() for role_id, role_data in role_dict.items():
if not any(role_data for role_data in role_dict.values()): # No roles # Skip non-configured roles
if not role_data:
continue continue
async for member in AsyncIter(guild.members): mr_dict = await self.config.custom("RoleMember", role_id, member.id).all()
has_roles = [r.id for r in member.roles]
# Stop if they've had the role and reapplying is disabled
if not reapply and mr_dict["had_role"]:
log.debug(f"{member.display_name} - Not reapplying")
continue
# Stop if the check_again_time hasn't passed yet
if (
mr_dict["check_again_time"] is not None
and datetime.fromisoformat(mr_dict["check_again_time"]) >= utcnow
):
log.debug(f"{member.display_name} - Not time to check again yet")
continue
member: discord.Member
has_roles = set(r.id for r in member.roles)
# Stop if they currently have or don't have the role, and mark had_role
if (role_id in has_roles and not role_data["remove"]) or (
role_id not in has_roles and role_data["remove"]
):
if not mr_dict["had_role"]:
await self.config.custom(
"RoleMember", role_id, member.id
).had_role.set(True)
log.debug(f"{member.display_name} - applying had_role")
continue
# Stop if they don't have all the required roles
if role_data is None or (
"required" in role_data and not set(role_data["required"]) & has_roles
):
continue
check_time = member.joined_at + timedelta(
days=role_data["days"],
hours=role_data.get("hours", 0),
)
# Check if enough time has passed to get the role and save the check_again_time
if check_time >= utcnow:
await self.config.custom(
"RoleMember", role_id, member.id
).check_again_time.set(check_time.isoformat())
log.debug(
f"{member.display_name} - Not enough time has passed to qualify for the role\n"
f"Waiting until {check_time}"
)
continue
if role_data["remove"]:
removelist.append(role_id)
else:
addlist.append(role_id)
# Done iterating through roles, now add or remove the roles
if not addlist and not removelist:
continue
# log.debug(f"{addlist=}\n{removelist=}")
add_roles = [ add_roles = [
int(rID) discord.utils.get(guild.roles, id=int(role_id)) for role_id in addlist
for rID, r_data in role_dict.items()
if r_data is not None and not r_data["remove"]
] ]
remove_roles = [ remove_roles = [
int(rID) discord.utils.get(guild.roles, id=int(role_id)) for role_id in removelist
for rID, r_data in role_dict.items()
if r_data is not None and r_data["remove"]
] ]
check_add_roles = set(add_roles) - set(has_roles) if None in add_roles or None in remove_roles:
check_remove_roles = set(remove_roles) & set(has_roles) log.info(
f"Timerole ran into an error with the roles in: {add_roles + remove_roles}"
)
await self.check_required_and_date( if addlist:
addlist, check_add_roles, has_roles, member, role_dict try:
await member.add_roles(*add_roles, reason="Timerole", atomic=False)
except (discord.Forbidden, discord.NotFound) as e:
log.exception("Failed Adding Roles")
add_results += f"{member.display_name} : **(Failed Adding Roles)**\n"
else:
add_results += " \n".join(
f"{member.display_name} : {role.name}" for role in add_roles
) )
await self.check_required_and_date( for role_id in addlist:
removelist, check_remove_roles, has_roles, member, role_dict await self.config.custom(
"RoleMember", role_id, member.id
).had_role.set(True)
if removelist:
try:
await member.remove_roles(*remove_roles, reason="Timerole", atomic=False)
except (discord.Forbidden, discord.NotFound) as e:
log.exception("Failed Removing Roles")
remove_results += f"{member.display_name} : **(Failed Removing Roles)**\n"
else:
remove_results += " \n".join(
f"{member.display_name} : {role.name}" for role in remove_roles
) )
for role_id in removelist:
await self.config.custom(
"RoleMember", role_id, member.id
).had_role.set(True)
# Done iterating through members, now maybe announce to the guild
channel = await self.config.guild(guild).announce() channel = await self.config.guild(guild).announce()
if channel is not None: if channel is not None:
channel = guild.get_channel(channel) channel = guild.get_channel(channel)
if add_results:
title = "**These members have received the following roles**\n" title = "**These members have received the following roles**\n"
await self.announce_roles(title, addlist, channel, guild, to_add=True) await announce_to_channel(channel, add_results, title)
if remove_results:
title = "**These members have lost the following roles**\n" title = "**These members have lost the following roles**\n"
await self.announce_roles(title, removelist, channel, guild, to_add=False) await announce_to_channel(channel, remove_results, title)
# End
async def announce_roles(self, title, role_list, channel, guild, to_add: True):
results = "" # async def announce_roles(self, title, role_list, channel, guild, to_add: True):
async for member, role_id in AsyncIter(role_list): # results = ""
role = discord.utils.get(guild.roles, id=role_id) # async for member, role_id in AsyncIter(role_list):
try: # role = discord.utils.get(guild.roles, id=role_id)
if to_add: # try:
await member.add_roles(role, reason="Timerole") # if to_add:
else: # await member.add_roles(role, reason="Timerole")
await member.remove_roles(role, reason="Timerole") # else:
except (discord.Forbidden, discord.NotFound) as e: # await member.remove_roles(role, reason="Timerole")
results += "{} : {} **(Failed)**\n".format(member.display_name, role.name) # except (discord.Forbidden, discord.NotFound) as e:
else: # results += f"{member.display_name} : {role.name} **(Failed)**\n"
results += "{} : {}\n".format(member.display_name, role.name) # else:
if channel is not None and results: # results += f"{member.display_name} : {role.name}\n"
await channel.send(title) # if channel is not None and results:
for page in pagify(results, shorten_by=50): # await channel.send(title)
await channel.send(page) # for page in pagify(results, shorten_by=50):
elif results: # Channel is None, log the results # await channel.send(page)
log.info(results) # elif results: # Channel is None, log the results
# log.info(results)
async def check_required_and_date(self, role_list, check_roles, has_roles, member, role_dict):
async for role_id in AsyncIter(check_roles): # async def check_required_and_date(self, role_list, check_roles, has_roles, member, role_dict):
# Check for required role # async for role_id in AsyncIter(check_roles):
if "required" in role_dict[str(role_id)]: # # Check for required role
if not set(role_dict[str(role_id)]["required"]) & set(has_roles): # if "required" in role_dict[str(role_id)]:
# Doesn't have required role # if not set(role_dict[str(role_id)]["required"]) & set(has_roles):
continue # # Doesn't have required role
# continue
if ( #
member.joined_at # if (
+ timedelta( # member.joined_at
days=role_dict[str(role_id)]["days"], # + timedelta(
hours=role_dict[str(role_id)].get("hours", 0), # days=role_dict[str(role_id)]["days"],
) # hours=role_dict[str(role_id)].get("hours", 0),
<= datetime.today() # )
): # <= datetime.utcnow()
# Qualifies # ):
role_list.append((member, role_id)) # # Qualifies
# role_list.append((member, role_id))
async def check_hour(self): async def check_hour(self):
await sleep_till_next_hour() await sleep_till_next_hour()

Loading…
Cancel
Save