import asyncio import logging import os from shlex import join from tempfile import TemporaryDirectory import discord import discord_tools.message import discord_tools.pixiv_link_extractor import pixiv.pixiv_api as pixiv_api logging.basicConfig( format="[%(asctime)s] [%(levelname)8s] %(name)s: %(message)s", level=logging.INFO ) REFRESH_TOKEN = os.getenv("PIXIV_REFRESH_TOKEN") if REFRESH_TOKEN is None: logging.fatal("Pixiv refresh token is missing, please set PIXIV_REFRESH_TOKEN") exit(1) DISCORD_TOKEN = os.getenv("DISCORD_TOKEN") if DISCORD_TOKEN is None: logging.fatal("Discord token is missing, please set DISCORD_TOKEN") exit(1) intents = discord.Intents.default() intents.message_content = True 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 async def on_ready(): logging.info(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: loop = asyncio.get_running_loop() data = await loop.run_in_executor( None, pixiv_api.download_pixiv_images, int(match.group("id")), tmp_dir, REFRESH_TOKEN, ) if data is None: await interaction.response.send_message( "Fail to extract pixiv image", ephemeral=True ) return (title, author) = data files = os.listdir(tmp_dir) files.sort() await discord_tools.message.send_command_reply( interaction, files, tmp_dir, True, title=title, art_author=author, 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: loop = asyncio.get_running_loop() data = await loop.run_in_executor( None, pixiv_api.download_pixiv_images, int(match.group("id")), tmp_dir, REFRESH_TOKEN, ) if data is None: await interaction.response.send_message( "Fail to extract pixiv image", ephemeral=True ) return files = os.listdir(tmp_dir) files.sort() (title, author) = data await discord_tools.message.send_command_reply( interaction=interaction, files=files, tmp_dir=tmp_dir, muted=muted, title=title, art_author=author, url=match.group("url"), author=message.author, ) 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): logging.info(f"Joined new guild: {guild.name}") @client.event async def on_message(message: discord.Message): if message.author == client.user: return (result, muted, match) = discord_tools.pixiv_link_extractor.should_process(message) if result and match is not None: with TemporaryDirectory() as tmp_dir: async with message.channel.typing(): loop = asyncio.get_running_loop() data = await loop.run_in_executor( None, pixiv_api.download_pixiv_images, int(match.group("id")), tmp_dir, REFRESH_TOKEN, ) if data is None: return (title, author) = data files = os.listdir(tmp_dir) files.sort() # await send_message.send_message_with_embed( # message, files, title, tmp_dir, match.group("url") # ) await discord_tools.message.send_message( message, files, tmp_dir, muted, title, author ) client.run(DISCORD_TOKEN)