add test command

master
bakatrouble 1 year ago
parent 728373cdd5
commit d10f9a52be

@ -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]

@ -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]
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)
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:

Loading…
Cancel
Save