From d10f9a52beb1a1e6744e4befd85bf48233d90fb8 Mon Sep 17 00:00:00 2001 From: bakatrouble Date: Thu, 1 Jun 2023 00:04:20 +0300 Subject: [PATCH] add test command --- e621.py | 3 ++ main.py | 149 ++++++++++++++++++++++++++++++++------------------------ 2 files changed, 87 insertions(+), 65 deletions(-) diff --git a/e621.py b/e621.py index 2f085ed..b99f2da 100644 --- a/e621.py +++ b/e621.py @@ -137,3 +137,6 @@ class E621: async def get_posts(self, tags='', page=1, limit=50) -> List[E621Post]: r = (await self.client.get('/posts.json', params={'tags': tags, 'page': page, 'limit': limit})).json() return [E621Post.from_dict(p) for p in r['posts']] + + async def get_post(self, post_id: str) -> E621Post: + return (await self.get_posts(f'id:{post_id}'))[0] diff --git a/main.py b/main.py index 22c4825..e8c03a6 100644 --- a/main.py +++ b/main.py @@ -2,17 +2,19 @@ import asyncio import logging import os from io import BytesIO +from typing import List + from PIL import Image import httpx import redis.asyncio as aioredis from aiogram import Bot, Dispatcher from aiogram.dispatcher import filters -from aiogram.types import Message, ParseMode +from aiogram.types import Message, ParseMode, ChatActions from aiogram.utils import executor, exceptions import dotenv -from e621 import E621 +from e621 import E621, E621Post dotenv.load_dotenv('.env') @@ -26,6 +28,72 @@ dp = Dispatcher(bot) user_ids = list(map(int, os.environ['USERS'].split(','))) +async def send_post(post: E621Post, tag_list: List[str]): + await bot.send_chat_action(int(os.environ['SEND_CHAT']), action=ChatActions.TYPING) + monitored_tags = set(post.tags.flatten()) & set(tag_list) + artist_tags = post.tags.artist + character_tags = post.tags.character + copyright_tags = post.tags.copyright + caption = '\n'.join(l for l in [ + f'Monitored tags: {" ".join(monitored_tags) or "None"}', + artist_tags and f'Artist: {" ".join(artist_tags)}', + character_tags and f'Character: {", ".join(character_tags)}', + copyright_tags and f'Copyright: {", ".join(copyright_tags)}', + f'\nhttps://e621.net/posts/{post.id}' + ] if l) + if post.file.url: + try: + logging.warning(post.file.url) + async with httpx.AsyncClient() as client: + file = BytesIO() + file.write((await client.get(post.file.url)).content) + file.name = f'file.{post.file.ext}' + file.seek(0) + if post.file.ext == 'webm': + if post.file.size < 50_000_000: + await bot.send_video(int(os.environ['SEND_CHAT']), + file, + width=post.file.width, + height=post.file.height, + thumb=post.preview.url, + caption=caption, + parse_mode=ParseMode.HTML) + else: + await bot.send_message(int(os.environ['SEND_CHAT']), + f'File is too large: {post.file.url}\n\n' + caption, + parse_mode=ParseMode.HTML) + elif post.file.ext == 'gif': + await bot.send_animation(int(os.environ['SEND_CHAT']), + file, + width=post.file.width, + height=post.file.height, + thumb=post.preview.url, + caption=caption, + parse_mode=ParseMode.HTML) + elif post.file.ext in ('png', 'jpg'): + if post.file.size > 10000000: + logging.warning('compressing') + dl_im = Image.open(file).convert('RGBA') + if dl_im.size[0] > 2000 or dl_im.size[1] > 2000: + larger_dimension = max(dl_im.size) + ratio = 2000 / larger_dimension + dl_im = dl_im.resize((int(dl_im.size[0] * ratio), int(dl_im.size[1] * ratio)), + Image.LANCZOS) + im = Image.new('RGBA', dl_im.size, (255, 255, 255)) + composite = Image.alpha_composite(im, dl_im).convert('RGB') + file = BytesIO() + composite.save(file, format='JPEG') + file.seek(0) + file.name = 'file.jpg' + await bot.send_photo(int(os.environ['SEND_CHAT']), + file, + caption=caption, + parse_mode=ParseMode.HTML) + await redis.sadd('e621:sent', post.id) + except exceptions.TelegramAPIError as e: + logging.exception(e) + + async def check_updates(): logging.warning('Waiting for lock...') async with redis.lock('e621:update'): @@ -42,69 +110,7 @@ async def check_updates(): for i in range(len(posts)): if already_sent[i]: continue - post = posts[i] - monitored_tags = set(post.tags.flatten()) & set(tag_list) - artist_tags = post.tags.artist - character_tags = post.tags.character - copyright_tags = post.tags.copyright - caption = '\n'.join(l for l in [ - f'Monitored tags: {" ".join(monitored_tags)}', - artist_tags and f'Artist: {" ".join(artist_tags)}', - character_tags and f'Character: {", ".join(character_tags)}', - copyright_tags and f'Copyright: {", ".join(copyright_tags)}', - f'\nhttps://e621.net/posts/{post.id}' - ] if l) - if post.file.url: - try: - logging.warning(post.file.url) - async with httpx.AsyncClient() as client: - file = BytesIO() - file.write((await client.get(post.file.url)).content) - file.name = f'file.{post.file.ext}' - file.seek(0) - if post.file.ext == 'webm': - if post.file.size < 50_000_000: - await bot.send_video(int(os.environ['SEND_CHAT']), - file, - width=post.file.width, - height=post.file.height, - thumb=post.preview.url, - caption=caption, - parse_mode=ParseMode.HTML) - else: - await bot.send_message(int(os.environ['SEND_CHAT']), - f'File is too large: {post.file.url}\n\n' + caption, - parse_mode=ParseMode.HTML) - elif post.file.ext == 'gif': - await bot.send_animation(int(os.environ['SEND_CHAT']), - file, - width=post.file.width, - height=post.file.height, - thumb=post.preview.url, - caption=caption, - parse_mode=ParseMode.HTML) - elif post.file.ext in ('png', 'jpg'): - if post.file.size > 10000000: - logging.warning('compressing') - dl_im = Image.open(file).convert('RGBA') - if dl_im.size[0] > 2000 or dl_im.size[1] > 2000: - larger_dimension = max(dl_im.size) - ratio = 2000 / larger_dimension - dl_im = dl_im.resize((int(dl_im.size[0] * ratio), int(dl_im.size[1] * ratio)), - Image.LANCZOS) - im = Image.new('RGBA', dl_im.size, (255, 255, 255)) - composite = Image.alpha_composite(im, dl_im).convert('RGB') - file = BytesIO() - composite.save(file, format='JPEG') - file.seek(0) - file.name = 'file.jpg' - await bot.send_photo(int(os.environ['SEND_CHAT']), - file, - caption=caption, - parse_mode=ParseMode.HTML) - await redis.sadd('e621:sent', post.id) - except exceptions.TelegramAPIError as e: - logging.exception(e) + await send_post(posts[i], tag_list) @dp.message_handler(filters.IDFilter(chat_id=user_ids), commands=['add']) @@ -168,6 +174,19 @@ async def update(msg: Message): await check_updates() +@dp.message_handler(filters.IDFilter(chat_id=user_ids), commands=['test']) +async def test(msg: Message): + args = msg.get_args() + if not args: + await msg.reply('Please provide post id') + return + post = await e621.get_posts(f'id:{args[0]}') + if not post: + await msg.reply('Post not found') + return + await send_post(post[0], []) + + async def background_on_start(): await redis.delete('e621:update') while True: