Migrate rolegiver, use components for add/remove roles, closes #101

This commit is contained in:
Zeva Rose 2022-02-03 15:46:17 -07:00
parent 68d29d9dd8
commit 0c382986cf

View file

@ -1,44 +1,38 @@
"""J.A.R.V.I.S. Role Giver Cog.""" """J.A.R.V.I.S. Role Giver Cog."""
import asyncio import asyncio
from discord import Role from dis_snek import InteractionContext, Permissions, Scale, Snake
from discord.ext import commands from dis_snek.client.utils import get
from discord_slash import SlashContext, cog_ext from dis_snek.models.discord.components import ActionRow, Select, SelectOption
from discord_slash.utils.manage_commands import create_option from dis_snek.models.discord.embed import EmbedField
from discord_slash.utils.manage_components import ( from dis_snek.models.discord.role import Role
create_actionrow, from dis_snek.models.snek.application_commands import (
create_select, OptionTypes,
create_select_option, slash_command,
wait_for_component, slash_option,
) )
from dis_snek.models.snek.command import cooldown
from dis_snek.models.snek.cooldowns import Buckets
from jarvis.db.models import Rolegiver from jarvis.db.models import Rolegiver
from jarvis.utils import build_embed from jarvis.utils import build_embed
from jarvis.utils.field import Field
from jarvis.utils.permissions import admin_or_permissions from jarvis.utils.permissions import admin_or_permissions
class RolegiverCog(commands.Cog): class RolegiverCog(Scale):
"""J.A.R.V.I.S. Role Giver Cog.""" """J.A.R.V.I.S. Role Giver Cog."""
def __init__(self, bot: commands.Bot): def __init__(self, bot: Snake):
self.bot = bot self.bot = bot
@cog_ext.cog_subcommand( @slash_command(
base="rolegiver", name="rolegiver", sub_cmd_name="add", sub_cmd_description="Add a role to rolegiver"
name="add",
description="Add a role to rolegiver",
options=[
create_option(
name="role",
description="Role to add",
option_type=8,
required=True,
) )
], @slash_option(
name="role", description="Role to add", optin_type=OptionTypes.ROLE, required=True
) )
@admin_or_permissions(manage_guild=True) @admin_or_permissions(Permissions.MANAGE_GUILD)
async def _rolegiver_add(self, ctx: SlashContext, role: Role) -> None: async def _rolegiver_add(self, ctx: InteractionContext, role: Role) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first() setting = Rolegiver.objects(guild=ctx.guild.id).first()
if setting and role.id in setting.roles: if setting and role.id in setting.roles:
await ctx.send("Role already in rolegiver", hidden=True) await ctx.send("Role already in rolegiver", hidden=True)
@ -58,7 +52,7 @@ class RolegiverCog(commands.Cog):
for role_id in setting.roles: for role_id in setting.roles:
if role_id == role.id: if role_id == role.id:
continue continue
e_role = ctx.guild.get_role(role_id) e_role = await ctx.guild.get_role(role_id)
if not e_role: if not e_role:
continue continue
roles.append(e_role) roles.append(e_role)
@ -67,8 +61,8 @@ class RolegiverCog(commands.Cog):
value = "\n".join([r.mention for r in roles]) if roles else "None" value = "\n".join([r.mention for r in roles]) if roles else "None"
fields = [ fields = [
Field(name="New Role", value=f"{role.mention}"), EmbedField(name="New Role", value=f"{role.mention}"),
Field(name="Existing Role(s)", value=value), EmbedField(name="Existing Role(s)", value=value),
] ]
embed = build_embed( embed = build_embed(
@ -79,21 +73,19 @@ class RolegiverCog(commands.Cog):
embed.set_thumbnail(url=ctx.guild.icon_url) embed.set_thumbnail(url=ctx.guild.icon_url)
embed.set_author( embed.set_author(
name=ctx.author.nick if ctx.author.nick else ctx.author.name, name=ctx.author.display_name,
icon_url=ctx.author.avatar_url, icon_url=ctx.author.display_avatar,
) )
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") embed.set_footer(text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}")
await ctx.send(embed=embed) await ctx.send(embed=embed)
@cog_ext.cog_subcommand( @slash_command(
base="rolegiver", name="rolegiver", sub_cmd_name="remove", sub_cmd_description="Remove a role from rolegiver"
name="remove",
description="Remove a role from rolegiver",
) )
@admin_or_permissions(manage_guild=True) @admin_or_permissions(Permissions.MANAGE_GUILD)
async def _rolegiver_remove(self, ctx: SlashContext) -> None: async def _rolegiver_remove(self, ctx: InteractionContext) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first() setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or (setting and not setting.roles): if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", hidden=True) await ctx.send("Rolegiver has no roles", hidden=True)
@ -101,28 +93,28 @@ class RolegiverCog(commands.Cog):
options = [] options = []
for role in setting.roles: for role in setting.roles:
role: Role = ctx.guild.get_role(role) role: Role = await ctx.guild.get_role(role)
option = create_select_option(label=role.name, value=str(role.id)) option = SelectOption(label=role.name, value=str(role.id))
options.append(option) options.append(option)
select = create_select( select = Select(
options=options, options=options,
custom_id="to_delete", custom_id="to_delete",
placeholder="Select roles to remove", placeholder="Select roles to remove",
min_values=1, min_values=1,
max_values=len(options), max_values=len(options),
) )
components = [create_actionrow(select)] components = [ActionRow(select)]
message = await ctx.send(content="\u200b", components=components) message = await ctx.send(content="\u200b", components=components)
try: try:
context = await wait_for_component( context = await self.bot.wait_for_component(
self.bot, self.bot,
check=lambda x: ctx.author.id == x.author.id, check=lambda x: ctx.author.id == x.author.id,
message=message, message=message,
timeout=60 * 1, timeout=60 * 1,
) )
for to_delete in context.selected_options: for to_delete in context.context.values:
setting.roles.remove(int(to_delete)) setting.roles.remove(int(to_delete))
setting.save() setting.save()
for row in components: for row in components:
@ -131,7 +123,7 @@ class RolegiverCog(commands.Cog):
roles = [] roles = []
for role_id in setting.roles: for role_id in setting.roles:
e_role = ctx.guild.get_role(role_id) e_role = await ctx.guild.get_role(role_id)
if not e_role: if not e_role:
continue continue
roles.append(e_role) roles.append(e_role)
@ -141,8 +133,8 @@ class RolegiverCog(commands.Cog):
value = "\n".join([r.mention for r in roles]) if roles else "None" value = "\n".join([r.mention for r in roles]) if roles else "None"
fields = [ fields = [
Field(name="Removed Role", value=f"{role.mention}"), EmbedField(name="Removed Role", value=f"{role.mention}"),
Field(name="Remaining Role(s)", value=value), EmbedField(name="Remaining Role(s)", value=value),
] ]
embed = build_embed( embed = build_embed(
@ -153,13 +145,15 @@ class RolegiverCog(commands.Cog):
embed.set_thumbnail(url=ctx.guild.icon_url) embed.set_thumbnail(url=ctx.guild.icon_url)
embed.set_author( embed.set_author(
name=ctx.author.nick if ctx.author.nick else ctx.author.name, name=ctx.author.display_name,
icon_url=ctx.author.avatar_url, icon_url=ctx.author.display_avatar,
) )
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") embed.set_footer(
text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}"
)
await context.edit_origin( await context.context.edit_origin(
content=f"Removed {len(context.selected_options)} role(s)", content=f"Removed {len(context.selected_options)} role(s)",
embed=embed, embed=embed,
components=components, components=components,
@ -170,12 +164,8 @@ class RolegiverCog(commands.Cog):
component["disabled"] = True component["disabled"] = True
await message.edit(components=components) await message.edit(components=components)
@cog_ext.cog_subcommand( @slash_command(name="rolegiver", sub_cmd_name="list", description="List rolegiver roles")
base="rolegiver", async def _rolegiver_list(self, ctx: InteractionContext) -> None:
name="list",
description="List roles rolegiver",
)
async def _rolegiver_list(self, ctx: SlashContext) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first() setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or (setting and not setting.roles): if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", hidden=True) await ctx.send("Rolegiver has no roles", hidden=True)
@ -183,7 +173,7 @@ class RolegiverCog(commands.Cog):
roles = [] roles = []
for role_id in setting.roles: for role_id in setting.roles:
e_role = ctx.guild.get_role(role_id) e_role = await ctx.guild.get_role(role_id)
if not e_role: if not e_role:
continue continue
roles.append(e_role) roles.append(e_role)
@ -201,21 +191,17 @@ class RolegiverCog(commands.Cog):
embed.set_thumbnail(url=ctx.guild.icon_url) embed.set_thumbnail(url=ctx.guild.icon_url)
embed.set_author( embed.set_author(
name=ctx.author.nick if ctx.author.nick else ctx.author.name, name=ctx.author.nick if ctx.author.nick else ctx.author.username,
icon_url=ctx.author.avatar_url, icon_url=ctx.author.display_avatar,
) )
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") embed.set_footer(text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}")
await ctx.send(embed=embed) await ctx.send(embed=embed)
@cog_ext.cog_subcommand( @slash_command(name="role", sub_cmd_name="get", sub_cmd_description="Get a role")
base="role", @cooldown(bucket=Buckets.USER, rate=1, interval=10)
name="get", async def _role_get(self, ctx: InteractionContext) -> None:
description="Get a role from rolegiver",
)
@commands.cooldown(1, 10, commands.BucketType.user)
async def _role_get(self, ctx: SlashContext) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first() setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or (setting and not setting.roles): if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", hidden=True) await ctx.send("Rolegiver has no roles", hidden=True)
@ -223,32 +209,44 @@ class RolegiverCog(commands.Cog):
options = [] options = []
for role in setting.roles: for role in setting.roles:
role: Role = ctx.guild.get_role(role) role: Role = await ctx.guild.get_role(role)
option = create_select_option(label=role.name, value=str(role.id)) option = SelectOption(label=role.name, value=str(role.id))
options.append(option) options.append(option)
select = create_select( select = Select(
options=options, options=options,
custom_id="to_delete", custom_id="to_delete",
placeholder="Select roles to remove", placeholder="Select roles to remove",
min_values=1, min_values=1,
max_values=len(options), max_values=len(options),
) )
components = [create_actionrow(select)] components = [ActionRow(select)]
_ = await ctx.send(content="\u200b", components=components) message = await ctx.send(content="\u200b", components=components)
await ctx.author.add_roles(role, reason="Rolegiver") try:
context = await self.bot.wait_for_component(
check=lambda x: ctx.author.id == x.author_id,
messages=message,
timeout=60 * 5,
)
added_roles = []
for role in context.context.values:
role = await ctx.guild.get_role(int(role))
added_roles.append(role)
await ctx.author.add_role(role, reason="Rolegiver")
roles = ctx.author.roles roles = ctx.author.roles
if roles: if roles:
roles.sort(key=lambda x: -x.position) roles.sort(key=lambda x: -x.position)
_ = roles.pop(-1) _ = roles.pop(-1)
avalue = "\n".join([r.mention for r in added_roles]) if added_roles else "None"
value = "\n".join([r.mention for r in roles]) if roles else "None" value = "\n".join([r.mention for r in roles]) if roles else "None"
fields = [ fields = [
Field(name="Added Role", value=f"{role.mention}"), EmbedField(name="Added Role(s)", value=avalue),
Field(name="Prior Role(s)", value=value), EmbedField(name="Prior Role(s)", value=value),
] ]
embed = build_embed( embed = build_embed(
@ -259,51 +257,77 @@ class RolegiverCog(commands.Cog):
embed.set_thumbnail(url=ctx.guild.icon_url) embed.set_thumbnail(url=ctx.guild.icon_url)
embed.set_author( embed.set_author(
name=ctx.author.nick if ctx.author.nick else ctx.author.name, name=ctx.author.display_name,
icon_url=ctx.author.avatar_url, icon_url=ctx.author.display_avatar,
) )
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") embed.set_footer(
text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}"
await ctx.send(embed=embed)
@cog_ext.cog_subcommand(
base="role",
name="forfeit",
description="Have rolegiver take away role",
options=[
create_option(
name="role",
description="Role to remove",
option_type=8,
required=True,
) )
],
) for row in components:
@commands.cooldown(1, 10, commands.BucketType.user) for component in row["components"]:
async def _role_forfeit(self, ctx: SlashContext, role: Role) -> None: component["disabled"] = True
await context.context.edit_origin(embed=embed, content="\u200b", components=components)
except asyncio.TimeoutError:
for row in components:
for component in row["components"]:
component["disabled"] = True
await message.edit(components=components)
@slash_command(name="role", sub_cmd_name="remove", sub_cmd_description="Remove a role")
@cooldown(bucket=Buckets.USER, rate=1, interval=10)
async def _role_remove(self, ctx: InteractionContext) -> None:
user_roles = ctx.author.roles
setting = Rolegiver.objects(guild=ctx.guild.id).first() setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or (setting and not setting.roles): if not setting or (setting and not setting.roles):
await ctx.send("Rolegiver has no roles", hidden=True) await ctx.send("Rolegiver has no roles", hidden=True)
return return
elif role.id not in setting.roles: elif not any(x.id in setting.roles for x in user_roles):
await ctx.send("Role not in rolegiver", hidden=True) await ctx.send("You have no rolegiver roles", hidden=True)
return
elif role not in ctx.author.roles:
await ctx.send("You do not have that role", hidden=True)
return return
await ctx.author.remove_roles(role, reason="Rolegiver") valid = list(filter(lambda x: x.id in setting.roles, user_roles))
options = []
for role in valid:
option = SelectOption(label=role.name, value=str(role.id))
options.append(option)
roles = ctx.author.roles select = Select(
if roles: options=options,
roles.sort(key=lambda x: -x.position) custom_id="to_remove",
_ = roles.pop(-1) placeholder="Select roles to remove",
min_values=1,
max_values=len(options),
)
components = [ActionRow(select)]
value = "\n".join([r.mention for r in roles]) if roles else "None" message = await ctx.send(content="\u200b", components=components)
try:
context = await self.bot.wait_for_component(
check=lambda x: ctx.author.id == x.author_id,
messages=message,
timeout=60 * 5,
)
removed_roles = []
for to_remove in context.context.values:
role = get(user_roles, id=int(to_remove))
await ctx.author.remove_role(role, reason="Rolegiver")
user_roles.remove(role)
removed_roles.append(role)
user_roles.sort(key=lambda x: -x.position)
_ = user_roles.pop(-1)
value = "\n".join([r.mention for r in user_roles]) if user_roles else "None"
rvalue = "\n".join([r.mention for r in removed_roles]) if removed_roles else "None"
fields = [ fields = [
Field(name="Taken Role", value=f"{role.mention}"), EmbedField(name="Removed Role(s)", value=rvalue),
Field(name="Remaining Role(s)", value=value), EmbedField(name="Remaining Role(s)", value=value),
] ]
embed = build_embed( embed = build_embed(
@ -314,26 +338,35 @@ class RolegiverCog(commands.Cog):
embed.set_thumbnail(url=ctx.guild.icon_url) embed.set_thumbnail(url=ctx.guild.icon_url)
embed.set_author( embed.set_author(
name=ctx.author.nick if ctx.author.nick else ctx.author.name, name=ctx.author.display_name,
icon_url=ctx.author.avatar_url, icon_url=ctx.author.display_avatar,
) )
embed.set_footer(text=f"{ctx.author.name}#{ctx.author.discriminator} | {ctx.author.id}") embed.set_footer(
text=f"{ctx.author.username}#{ctx.author.discriminator} | {ctx.author.id}"
await ctx.send(embed=embed)
@cog_ext.cog_subcommand(
base="rolegiver",
name="cleanup",
description="Cleanup rolegiver roles",
) )
@admin_or_permissions(manage_guild=True)
async def _rolegiver_cleanup(self, ctx: SlashContext) -> None: for row in components:
for component in row["components"]:
component["disabled"] = True
await context.context.edit_origin(embed=embed, components=components, content="\u200b")
except asyncio.TimeoutError:
for row in components:
for component in row["components"]:
component["disabled"] = True
await message.edit(components=components)
@slash_command(
name="rolegiver", sub_cmd_name="cleanup", description="Removed deleted roles from rolegiver"
)
@admin_or_permissions(Permissions.MANAGE_GUILD)
async def _rolegiver_cleanup(self, ctx: InteractionContext) -> None:
setting = Rolegiver.objects(guild=ctx.guild.id).first() setting = Rolegiver.objects(guild=ctx.guild.id).first()
if not setting or not setting.roles: if not setting or not setting.roles:
await ctx.send("Rolegiver has no roles", hidden=True) await ctx.send("Rolegiver has no roles", hidden=True)
guild_roles = await ctx.guild.fetch_roles() guild_role_ids = [r.id for r in ctx.guild.roles]
guild_role_ids = [x.id for x in guild_roles]
for role_id in setting.roles: for role_id in setting.roles:
if role_id not in guild_role_ids: if role_id not in guild_role_ids:
setting.roles.remove(role_id) setting.roles.remove(role_id)
@ -342,6 +375,7 @@ class RolegiverCog(commands.Cog):
await ctx.send("Rolegiver cleanup finished") await ctx.send("Rolegiver cleanup finished")
def setup(bot: commands.Bot) -> None: def setup(bot: Snake) -> None:
"""Add RolegiverCog to J.A.R.V.I.S.""" """Add RolegiverCog to J.A.R.V.I.S."""
bot.add_cog(RolegiverCog(bot)) bot.add_cog(RolegiverCog(bot))
bot.add_cog(RolegiverCog(bot))