bakatrouble
/
e621_bot
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

add test command

Browse Source
master
bakatrouble 12 months ago
parent 728373cdd5
commit d10f9a52be
2 changed files with 87 additions and 65 deletions
Show Stats Download Patch File Download Diff File

3
e621.py
Unescape View File

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

149
main.py
Unescape View File

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

Write Preview
Loading…
Cancel
Save