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: