pixivcord/main.py

import math
import os
import re
from shlex import join
from tempfile import TemporaryDirectory
import discord
import pixiv

REFRESH_TOKEN = os.getenv("PIXIV_REFRESH_TOKEN")
if REFRESH_TOKEN is None:
    print("Pixiv refresh token is missing, please set PIXIV_REFRESH_TOKEN")
    exit(1)
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
if DISCORD_TOKEN is None:
    print("Discord token is missing, please set DISCORD_TOKEN")
    exit(1)


intents = discord.Intents.default()
intents.message_content = True

client = discord.Client(intents=intents)


@client.event
async def on_ready():
    print(f"We have logged in as {client.user}")


@client.event
async def on_message(message: discord.Message):
    if message.author == client.user:
        return

    if (
        match := re.search(
            r"(?P<url>https?://(?:www\.)?pixiv.net/\w+/artworks/(?P<id>\d+))",
            message.content,
        )
    ) is not None:
        with TemporaryDirectory() as tmp_dir:
            async with message.channel.typing():
                title = pixiv.dowload_pixiv_images(
                    int(match.group("id")), tmp_dir, REFRESH_TOKEN
                )
                files = os.listdir(tmp_dir)
                to_send_files = []
                to_send_embed = []
                files.reverse()
                for file in files:
                    joined = os.path.join(tmp_dir, file)
                    embed = discord.Embed(url=match.group("url"))
                    embed.set_image(url=f"attachment://{file}")
                    to_send_files.append(discord.File(joined, file))
                    to_send_embed.append(embed)
                if len(to_send_embed) > 4:
                    parts = math.ceil(len(to_send_embed) / 4)
                    i = 1
                    while len(to_send_embed) > 4:
                        splited_file = to_send_files[:4]
                        to_send_files = to_send_files[4:]
                        splited_embed = to_send_embed[:4]
                        to_send_embed = to_send_embed[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_embed:
                        embed.title = f"{title} - Part {i}/{parts}"
                    await message.channel.send(
                        files=to_send_files, embeds=to_send_embed
                    )
                else:
                    for embed in to_send_embed:
                        embed.title = title
                    await message.reply(files=to_send_files, embeds=to_send_embed)


client.run(DISCORD_TOKEN)