diff --git a/bot.py b/bot.py index 8cd3c15..c1bc236 100644 --- a/bot.py +++ b/bot.py @@ -20,6 +20,7 @@ from discord.ext import commands from rich.logging import RichHandler +import src.mongo.models from env import env from src.discord.globals import ( CHANNEL_BOTSPAM, @@ -164,6 +165,7 @@ async def setup_hook(self) -> None: await init_beanie( database=self.mongo_database.client["data"], document_models=[ + src.mongo.models.Event, # TODO ], ) diff --git a/src/discord/globals.py b/src/discord/globals.py index b4bd50b..c422a3b 100644 --- a/src/discord/globals.py +++ b/src/discord/globals.py @@ -2,6 +2,7 @@ Holds global variables shared between cogs and variables that are initialized when the bot is first setup. """ +from src.mongo.models import Event ############## # CONSTANTS @@ -119,7 +120,7 @@ ############## fish_now = 0 CENSOR = {} -EVENT_INFO = [] +EVENT_INFO: list[Event] = [] PING_INFO = [] INVITATIONAL_INFO = [] REPORTS = [] diff --git a/src/discord/membercommands.py b/src/discord/membercommands.py index 814d802..2b138ca 100644 --- a/src/discord/membercommands.py +++ b/src/discord/membercommands.py @@ -1182,7 +1182,7 @@ async def events( event_ten, ] param_list = [p for p in param_list if p is not None] - event_names = [e["name"] for e in src.discord.globals.EVENT_INFO] + event_names = [e.name for e in src.discord.globals.EVENT_INFO] selected_roles = [ discord.utils.get(member.guild.roles, name=e) @@ -1246,9 +1246,9 @@ async def events_autocomplete( List[app_commands.Choice[str]]: A list of string choices to choose from. """ return [ - app_commands.Choice(name=e["name"], value=e["name"]) + app_commands.Choice(name=e.name, value=e.name) for e in src.discord.globals.EVENT_INFO - if current.lower() in e["name"].lower() + if current.lower() in e.name.lower() ][:25] @app_commands.command(description="Gets a tag.") diff --git a/src/discord/staff/events.py b/src/discord/staff/events.py index c796034..b26e35e 100644 --- a/src/discord/staff/events.py +++ b/src/discord/staff/events.py @@ -15,6 +15,7 @@ ROLE_STAFF, ROLE_VIP, ) +from src.mongo.models import Event if TYPE_CHECKING: from bot import PiBot @@ -55,7 +56,7 @@ async def event_add( ) # Check to see if event has already been added. - if event_name in [e["name"] for e in src.discord.globals.EVENT_INFO]: + if event_name in [e.name for e in src.discord.globals.EVENT_INFO]: return await interaction.edit_original_response( content=f"The `{event_name}` event has already been added.", ) @@ -65,11 +66,11 @@ async def event_add( aliases_array = [] if event_aliases: aliases_array = re.findall(r"\w+", event_aliases) - new_dict = {"name": event_name, "aliases": aliases_array} + new_dict = Event(name=event_name, aliases=aliases_array, emoji=None) # Add dict into events container + await new_dict.insert() src.discord.globals.EVENT_INFO.append(new_dict) - await self.bot.mongo_database.insert("data", "events", new_dict) # Create role on server server = self.bot.get_guild(env.server_id) @@ -110,7 +111,7 @@ async def event_remove( # Check to make sure event has previously been added event_not_in_list = event_name not in [ - e["name"] for e in src.discord.globals.EVENT_INFO + e.name for e in src.discord.globals.EVENT_INFO ] # Check to see if role exists on server @@ -138,8 +139,8 @@ async def event_remove( event = next( e for e in src.discord.globals.EVENT_INFO if e["name"] == event_name ) + await event.delete() src.discord.globals.EVENT_INFO.remove(event) - await self.bot.mongo_database.delete("data", "events", event["_id"]) # Notify staff member of completion if delete_role == "yes": diff --git a/src/discord/tasks.py b/src/discord/tasks.py index de54b91..86276fb 100644 --- a/src/discord/tasks.py +++ b/src/discord/tasks.py @@ -13,6 +13,7 @@ from env import env from src.discord.invitationals import update_invitational_list from src.discord.views import UnselfmuteView +from src.mongo.models import Event if TYPE_CHECKING: from bot import PiBot @@ -77,7 +78,7 @@ async def pull_prev_info(self): src.discord.globals.REPORTS = await self.bot.mongo_database.get_reports() src.discord.globals.PING_INFO = await self.bot.mongo_database.get_pings() src.discord.globals.TAGS = await self.bot.mongo_database.get_tags() - src.discord.globals.EVENT_INFO = await self.bot.mongo_database.get_events() + src.discord.globals.EVENT_INFO = await Event.find_all().to_list() self.bot.settings = await self.bot.mongo_database.get_settings() assert isinstance(self.bot.settings, dict) diff --git a/src/discord/welcome.py b/src/discord/welcome.py index aacc841..e913da0 100644 --- a/src/discord/welcome.py +++ b/src/discord/welcome.py @@ -293,7 +293,7 @@ async def request_access( event_options: list[discord.SelectOption] = [] for event in src.discord.globals.EVENT_INFO: event_options.append( - discord.SelectOption(label=event["name"], emoji=event["emoji"]), + discord.SelectOption(label=event.name, emoji=event.emoji), ) event_options.sort(key=lambda x: x.label) event_chooser = Chooser( diff --git a/src/mongo/models.py b/src/mongo/models.py index 64cd474..adeafb4 100644 --- a/src/mongo/models.py +++ b/src/mongo/models.py @@ -1,3 +1,15 @@ """ Contains all database models """ + +from beanie import Document + + +class Event(Document): + name: str + aliases: list[str] + emoji: str | None + + class Settings: + name = "events" + use_cache = True diff --git a/src/mongo/mongo.py b/src/mongo/mongo.py index eef3f2e..28ae930 100644 --- a/src/mongo/mongo.py +++ b/src/mongo/mongo.py @@ -119,12 +119,6 @@ async def get_reports(self): """ return await self.get_entire_collection("data", "reports") - async def get_events(self): - """ - Gets all documents in the events collection. - """ - return await self.get_entire_collection("data", "events") - async def get_settings(self): """ Gets the one document containing settings information from the settings