From 50034cf7b0a5615e17641af96df1bf20e0c0320e Mon Sep 17 00:00:00 2001 From: sclement Date: Tue, 31 Oct 2023 15:31:07 +0100 Subject: [PATCH] Add slash command --- src/discord_tools/message.py | 73 +++++++++++------------ src/discord_tools/pixiv_link_extractor.py | 7 +++ src/main.py | 71 +++++++++++++++++++++- src/pixiv/pixiv_api.py | 3 +- 4 files changed, 114 insertions(+), 40 deletions(-) diff --git a/src/discord_tools/message.py b/src/discord_tools/message.py index cfc2ab0..feaf0cd 100644 --- a/src/discord_tools/message.py +++ b/src/discord_tools/message.py @@ -3,43 +3,7 @@ import os import discord -TIPS_MESSAGE = ":bulb: TIPS: You can disable Discord's default preview for link by wrapping it around `<` and `>`." - - -async def send_message_with_embed( - message: discord.Message, files, title: str, tmp_dir: str, url: str -): - to_send_files = [] - to_send_embeds = [] - - for file in files: - joined = os.path.join(tmp_dir, file) - embed = discord.Embed(url=url) - embed.set_image(url=f"attachment://{file}") - to_send_files.append(discord.File(joined, file)) - to_send_embeds.append(embed) - if len(to_send_embeds) > 4: - parts = math.ceil(len(to_send_embeds) / 4) - i = 1 - while len(to_send_embeds) > 4: - splited_file = to_send_files[:4] - to_send_files = to_send_files[4:] - splited_embed = to_send_embeds[:4] - to_send_embeds = to_send_embeds[4:] - for embed in splited_embed: - embed.title = f"{title} - Part {i}/{parts}" - if i == 1: - await message.reply(files=splited_file, embeds=splited_embed) - else: - await message.channel.send(files=splited_file, embeds=splited_embed) - i = i + 1 - for embed in to_send_embeds: - embed.title = f"{title} - Part {i}/{parts}" - await message.channel.send(files=to_send_files, embeds=to_send_embeds) - else: - for embed in to_send_embeds: - embed.title = title - await message.reply(files=to_send_files, embeds=to_send_embeds) +TIPS_MESSAGE = ":bulb: TIPS: You can disable Discord's default preview for link by wrapping it around `<` `>`" async def send_message(message: discord.Message, files, tmp_dir: str, muted: bool): @@ -64,3 +28,38 @@ async def send_message(message: discord.Message, files, tmp_dir: str, muted: boo await message.reply( files=to_send_files, content=(TIPS_MESSAGE if not muted else None) ) + + +async def send_command_reply( + interaction: discord.Interaction, + files, + tmp_dir: str, + muted: bool, + title: str, + url: str, + author: discord.User | discord.Member, +): + to_send_files = [] + for file in files: + joined = os.path.join(tmp_dir, file) + to_send_files.append(discord.File(joined, file)) + if len(to_send_files) > 10: + if interaction.channel is not None: + first = True + while len(to_send_files) > 10: + splited_file = to_send_files[:10] + to_send_files = to_send_files[10:] + if first: + await interaction.followup.send( + files=splited_file, + content=f"{title} - Submited by {author.mention} - <{url}> {chr(10) + TIPS_MESSAGE if not muted else ''}", + ) + first = False + else: + await interaction.channel.send(files=splited_file) # type: ignore + await interaction.channel.send(files=to_send_files) # type: ignore + else: + await interaction.followup.send( + files=to_send_files, + content=f"{title} - Submited by {author.mention} - <{url}> {chr(10) + TIPS_MESSAGE if not muted else ''}", + ) diff --git a/src/discord_tools/pixiv_link_extractor.py b/src/discord_tools/pixiv_link_extractor.py index 1071603..15cb529 100644 --- a/src/discord_tools/pixiv_link_extractor.py +++ b/src/discord_tools/pixiv_link_extractor.py @@ -16,3 +16,10 @@ def should_process( else: return (False, False, match) return (False, False, None) + + +def check_pixiv_url(url: str): + if match := re.search(link_regexp, url): + return (True, match) + else: + return (False, None) diff --git a/src/main.py b/src/main.py index 763931d..badff95 100644 --- a/src/main.py +++ b/src/main.py @@ -21,7 +21,17 @@ if DISCORD_TOKEN is None: intents = discord.Intents.default() intents.message_content = True -client = discord.Client(intents=intents) + +class MyClient(discord.Client): + def __init__(self, *, intents: discord.Intents): + super().__init__(intents=intents) + self.tree = discord.app_commands.CommandTree(self) + + async def setup_hook(self): + await self.tree.sync() + + +client = MyClient(intents=intents) @client.event @@ -29,6 +39,65 @@ async def on_ready(): print(f"We have logged in as {client.user}") +@client.tree.command() +@discord.app_commands.describe( + url="Pixiv art url", +) +async def pixivprev(interaction: discord.Interaction, url: str): + """Create pixiv preview""" + (result, match) = discord_tools.pixiv_link_extractor.check_pixiv_url(url) + if result and match is not None: + await interaction.response.defer() + with TemporaryDirectory() as tmp_dir: + title = pixiv_api.dowload_pixiv_images( + int(match.group("id")), tmp_dir, REFRESH_TOKEN + ) + files = os.listdir(tmp_dir) + files.sort() + await discord_tools.message.send_command_reply( + interaction, + files, + tmp_dir, + True, + title=title, + url=url, + author=interaction.user, + ) + else: + await interaction.response.send_message("Invalid pixiv url", ephemeral=True) + + +@client.tree.context_menu(name="Create pixiv preview") +async def pixiv_preview(interaction: discord.Interaction, message: discord.Message): + (result, muted, match) = discord_tools.pixiv_link_extractor.should_process(message) + if result and match is not None: + await interaction.response.defer() + with TemporaryDirectory() as tmp_dir: + title = pixiv_api.dowload_pixiv_images( + int(match.group("id")), tmp_dir, REFRESH_TOKEN + ) + files = os.listdir(tmp_dir) + files.sort() + await discord_tools.message.send_command_reply( + interaction=interaction, + files=files, + tmp_dir=tmp_dir, + muted=muted, + title=title, + url=match.group("url"), + author=interaction.user, + ) + else: + await interaction.response.send_message( + "Can't find any valid pixiv link.", ephemeral=True + ) + + +@client.event +async def on_guild_join(guild: discord.Guild): + print(f"Joined new guild: {guild.name}") + + @client.event async def on_message(message: discord.Message): if message.author == client.user: diff --git a/src/pixiv/pixiv_api.py b/src/pixiv/pixiv_api.py index 2bfe2bd..cb6229a 100644 --- a/src/pixiv/pixiv_api.py +++ b/src/pixiv/pixiv_api.py @@ -13,10 +13,9 @@ def dowload_pixiv_images( illust = json_result.illust pages = illust["meta_pages"] if len(pages) > 0: - print(f"Dowloading mutiple images to {dest_folder}") + print(f"Dowloading mutiple images to {dest_folder} ({len(pages)})") for page in pages: url = page["image_urls"]["original"] - print(url) api.download(url=url, path=dest_folder) else: print(f"Dowloading single image to {dest_folder}")