add test command

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]

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: <b>{" ".join(monitored_tags) or "None"}</b>',
+        artist_tags and f'Artist: <b>{" ".join(artist_tags)}</b>',
+        character_tags and f'Character: <b>{", ".join(character_tags)}</b>',
+        copyright_tags and f'Copyright: <b>{", ".join(copyright_tags)}</b>',
+        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]
+                await send_post(posts[i], tag_list)
-                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: <b>{" ".join(monitored_tags)}</b>',
-                    artist_tags and f'Artist: <b>{" ".join(artist_tags)}</b>',
-                    character_tags and f'Character: <b>{", ".join(character_tags)}</b>',
-                    copyright_tags and f'Copyright: <b>{", ".join(copyright_tags)}</b>',
-                    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)
 
 
 @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: