diff --git a/redbot/cogs/trivia/trivia.py b/redbot/cogs/trivia/trivia.py index b288be60c..6ed9ded72 100644 --- a/redbot/cogs/trivia/trivia.py +++ b/redbot/cogs/trivia/trivia.py @@ -1,17 +1,19 @@ """Module for Trivia cog.""" +import asyncio import math import pathlib from collections import Counter from typing import List - +import io import yaml import discord - from redbot.core import commands from redbot.core import Config, checks from redbot.core.data_manager import cog_data_path from redbot.core.i18n import Translator, cog_i18n from redbot.core.utils.chat_formatting import box, pagify, bold +from redbot.core.utils.predicates import MessagePredicate, ReactionPredicate +from redbot.core.utils.menus import start_adding_reactions from redbot.cogs.bank import check_global_setting_admin from .log import LOG from .session import TriviaSession @@ -183,6 +185,86 @@ class Trivia(commands.Cog): else: await ctx.send(_("Done. I will no longer reward the winner with a payout.")) + @triviaset.group(name="custom") + @commands.is_owner() + async def triviaset_custom(self, ctx: commands.Context): + """Manage Custom Trivia lists.""" + pass + + @triviaset_custom.command(name="list") + async def custom_trivia_list(self, ctx: commands.Context): + """List uploaded custom trivia.""" + personal_lists = sorted([p.resolve().stem for p in cog_data_path(self).glob("*.yaml")]) + no_lists_uploaded = _("No custom Trivia lists uploaded.") + + if not personal_lists: + if await ctx.embed_requested(): + await ctx.send( + embed=discord.Embed( + colour=await ctx.embed_colour(), description=no_lists_uploaded + ) + ) + else: + await ctx.send(no_lists_uploaded) + return + + if await ctx.embed_requested(): + await ctx.send( + embed=discord.Embed( + title=_("Uploaded trivia lists"), + colour=await ctx.embed_colour(), + description=", ".join(sorted(personal_lists)), + ) + ) + else: + msg = box( + bold(_("Uploaded trivia lists")) + "\n\n" + ", ".join(sorted(personal_lists)) + ) + if len(msg) > 1000: + await ctx.author.send(msg) + else: + await ctx.send(msg) + + @commands.is_owner() + @triviaset_custom.command(name="upload", aliases=["add"]) + async def trivia_upload(self, ctx: commands.Context): + """Upload a trivia file.""" + if not ctx.message.attachments: + await ctx.send(_("Supply a file with next message or type anything to cancel.")) + try: + message = await ctx.bot.wait_for( + "message", check=MessagePredicate.same_context(ctx), timeout=30 + ) + except asyncio.TimeoutError: + await ctx.send(_("You took too long to upload a list.")) + return + if not message.attachments: + await ctx.send(_("You have cancelled the upload process.")) + return + parsedfile = message.attachments[0] + else: + parsedfile = ctx.message.attachments[0] + try: + await self._save_trivia_list(ctx=ctx, attachment=parsedfile) + except yaml.error.MarkedYAMLError as exc: + await ctx.send(_("Invalid syntax: ") + str(exc)) + except yaml.error.YAMLError: + await ctx.send( + _("There was an error parsing the trivia list. See logs for more info.") + ) + LOG.exception("Custom Trivia file %s failed to upload", parsedfile.filename) + + @commands.is_owner() + @triviaset_custom.command(name="delete", aliases=["remove"]) + async def trivia_delete(self, ctx: commands.Context, name: str): + """Delete a trivia file.""" + filepath = cog_data_path(self) / f"{name}.yaml" + if filepath.exists(): + filepath.unlink() + await ctx.send(_("Trivia {filename} was deleted.").format(filename=filepath.stem)) + else: + await ctx.send(_("Trivia file was not found.")) + @commands.group(invoke_without_command=True) @commands.guild_only() async def trivia(self, ctx: commands.Context, *categories: str): @@ -525,6 +607,73 @@ class Trivia(commands.Cog): else: return dict_ + async def _save_trivia_list( + self, ctx: commands.Context, attachment: discord.Attachment + ) -> None: + """Checks and saves a trivia list to data folder. + + Parameters + ---------- + file : discord.Attachment + A discord message attachment. + + Returns + ------- + None + """ + filename = attachment.filename.rsplit(".", 1)[0] + + # Check if trivia filename exists in core files or if it is a command + if filename in self.trivia.all_commands or any( + filename == item.stem for item in get_core_lists() + ): + await ctx.send( + _( + "{filename} is a reserved trivia name and cannot be replaced.\n" + "Choose another name." + ).format(filename=filename) + ) + return + + file = cog_data_path(self) / f"{filename}.yaml" + if file.exists(): + overwrite_message = _("{filename} already exists. Do you wish to overwrite?").format( + filename=filename + ) + + can_react = ctx.channel.permissions_for(ctx.me).add_reactions + if not can_react: + overwrite_message += " (y/n)" + + overwrite_message_object: discord.Message = await ctx.send(overwrite_message) + if can_react: + # noinspection PyAsyncCall + start_adding_reactions( + overwrite_message_object, ReactionPredicate.YES_OR_NO_EMOJIS + ) + pred = ReactionPredicate.yes_or_no(overwrite_message_object, ctx.author) + event = "reaction_add" + else: + pred = MessagePredicate.yes_or_no(ctx=ctx) + event = "message" + try: + await ctx.bot.wait_for(event, check=pred, timeout=30) + except asyncio.TimeoutError: + await ctx.send(_("You took too long answering.")) + return + + if pred.result is False: + await ctx.send(_("I am not replacing the existing file.")) + return + + buffer = io.BytesIO(await attachment.read()) + yaml.safe_load(buffer) + buffer.seek(0) + + with file.open("wb") as fp: + fp.write(buffer.read()) + await ctx.send(_("Saved Trivia list as {filename}.").format(filename=filename)) + def _get_trivia_session(self, channel: discord.TextChannel) -> TriviaSession: return next( (session for session in self.trivia_sessions if session.ctx.channel == channel), None