pixivcord/main.py
2023-10-30 14:58:04 +01:00

84 lines
2.8 KiB
Python

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():
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"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"Part {i}/{parts}"
await message.channel.send(
files=to_send_files, embeds=to_send_embed
)
client.run(DISCORD_TOKEN)