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

add api

Browse Source
This commit is contained in:
bakatrouble
2025-07-14 17:12:42 +03:00
parent 6029e1b16d
commit a289d24698
5 changed files with 296 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
main.py
View File

@@ -3,7 +3,6 @@ import base64
import datetime
import logging
import os
import random
import re
import traceback
from asyncio import sleep
@@ -20,12 +19,13 @@ from PIL import Image
import httpx
import redis.asyncio as aioredis
from aiogram import Bot, Dispatcher, filters, exceptions, F
from aiogram import Bot, Dispatcher, filters, F
from aiogram.enums import ChatAction, ParseMode
from aiogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton, BufferedInputFile, \
CallbackQuery
import dotenv
from const import REDIS_SUBS_KEY, REDIS_LAST_VERSION_KEY, REDIS_SENT_KEY, REDIS_LOCK_KEY
from e621 import E621, E621Post, E621PostFile, E621PostVersion
dotenv.load_dotenv('.env')
@@ -127,7 +127,7 @@ async def send_post(post: E621Post, tag_list: Iterable[Iterable[str]]):
caption=caption,
parse_mode=ParseMode.HTML,
reply_markup=markup)
await redis.sadd('e621:sent', post.id)
await redis.sadd(REDIS_SENT_KEY, post.id)
except Exception as e:
logging.exception(e)
except Exception as e:
@@ -136,12 +136,12 @@ async def send_post(post: E621Post, tag_list: Iterable[Iterable[str]]):
async def check_updates():
logging.warning('Waiting for lock...')
async with redis.lock('e621:update'):
async with redis.lock(REDIS_LOCK_KEY):
logging.warning('Lock acquired...')
matched_posts = []
tag_list = set(tuple(t.decode().split()) for t in await redis.smembers('e621:subs'))
tag_list = set(tuple(t.decode().split()) for t in await redis.smembers(REDIS_SUBS_KEY))
tag_list_flat = set(sum(tag_list, ()))
last_post_version = int((await redis.get('e621:last_version') or b'0').decode())
last_post_version = int((await redis.get(REDIS_LAST_VERSION_KEY) or b'0').decode())
post_versions: List[E621PostVersion] = []
logging.warning(f'Getting post versions from id {last_post_version}')
for page in count(1):
@@ -167,7 +167,7 @@ async def check_updates():
break
matched_posts.sort()
if matched_posts:
already_sent: List = await redis.smismember('e621:sent', matched_posts)
already_sent: List = await redis.smismember(REDIS_SENT_KEY, matched_posts)
posts_to_send = [post_id for post_id, sent in zip(matched_posts, already_sent) if not sent]
logging.warning(f'Found {len(posts_to_send)} posts')
for post_chunk_idx in range(0, len(posts_to_send), PAGE_SIZE):
@@ -176,9 +176,9 @@ async def check_updates():
for i, post in enumerate(posts):
logging.warning(f'Sending post {post_chunk_idx + i + 1}/{len(posts_to_send)}')
await send_post(post, tag_list)
await redis.sadd('e621:sent', post.id)
await redis.sadd(REDIS_SENT_KEY, post.id)
await sleep(1)
await redis.set('e621:last_version', last_post_version)
await redis.set(REDIS_LAST_VERSION_KEY, last_post_version)
@dp.message(filters.Command('resend_after'), ChatFilter)
@@ -190,8 +190,8 @@ async def resend_after(msg: Message):
await msg.reply('Invalid timestamp or not provided')
return
async with redis.lock('e621:update'):
tag_list = [tuple(t.decode().split()) for t in await redis.smembers('e621:subs')]
async with redis.lock(REDIS_LOCK_KEY):
tag_list = [tuple(t.decode().split()) for t in await redis.smembers(REDIS_SUBS_KEY)]
for i, tag in enumerate(tag_list):
await msg.reply(f'Checking tag <b>{tag}</b> ({i+1}/{len(tag_list)})', parse_mode=ParseMode.HTML)
posts = []
@@ -218,8 +218,8 @@ async def add_tag(msg: Message):
return
for tag in args.split():
posts = await e621.get_posts(tag)
await redis.sadd('e621:sent', *[post.id for post in posts])
await redis.sadd('e621:subs', tag)
await redis.sadd(REDIS_SENT_KEY, *[post.id for post in posts])
await redis.sadd(REDIS_SUBS_KEY, tag)
await msg.reply(f'Tags {args} added')
@@ -233,20 +233,20 @@ async def add_tags(msg: Message):
tags.sort()
tags = ' '.join(tags)
posts = await e621.get_posts(tags)
await redis.sadd('e621:sent', *[post.id for post in posts])
await redis.sadd('e621:subs', tags)
await redis.sadd(REDIS_SENT_KEY, *[post.id for post in posts])
await redis.sadd(REDIS_SUBS_KEY, tags)
await msg.reply(f'Tag group <code>{tags}</code> added', parse_mode=ParseMode.HTML)
@dp.message(filters.Command('mark_old_as_sent'), ChatFilter)
async def mark_old_as_sent(msg: Message):
logging.warning('Waiting for lock...')
async with redis.lock('e621:update'):
tag_list = [t.decode() for t in await redis.smembers('e621:subs')]
async with redis.lock(REDIS_LOCK_KEY):
tag_list = [t.decode() for t in await redis.smembers(REDIS_SUBS_KEY)]
m = await msg.reply(f'0/{len(tag_list)} tags have old posts marked as sent')
for i, tag in enumerate(tag_list, 1):
posts = await e621.get_posts(tag)
await redis.sadd('e621:sent', *[post.id for post in posts])
await redis.sadd(REDIS_SENT_KEY, *[post.id for post in posts])
await m.edit_text(f'{i}/{len(tag_list)} tags have old posts marked as sent')
await sleep(1)
await m.edit_text(f'Done marking old posts as sent for {len(tag_list)} tags')
@@ -261,10 +261,10 @@ async def del_tag(msg: Message):
if ' ' in args:
await msg.reply('Tag should not contain spaces')
return
if not await redis.sismember('e621:subs', args):
if not await redis.sismember(REDIS_SUBS_KEY, args):
await msg.reply('Tag not found')
return
await redis.srem('e621:subs', args)
await redis.srem(REDIS_SUBS_KEY, args)
await msg.reply(f'Tag {args} removed')
@@ -275,13 +275,13 @@ async def del_command(msg: Message):
await msg.reply('Please provide tag to subscribe to')
return
for tag in args.split():
await redis.srem('e621:subs', tag)
await redis.srem(REDIS_SUBS_KEY, tag)
await msg.reply(f'Tags {args} removed')
@dp.message(filters.Command('list'), ChatFilter)
async def list_tags(msg: Message):
tags = [t.decode() for t in await redis.smembers('e621:subs')]
tags = [t.decode() for t in await redis.smembers(REDIS_SUBS_KEY)]
tags.sort()
lines = []
for tag in tags:
@@ -313,7 +313,7 @@ async def test(msg: Message):
if not post:
await msg.reply('Post not found')
return
tag_list = [tuple(t.decode().split()) for t in await redis.smembers('e621:subs')]
tag_list = [tuple(t.decode().split()) for t in await redis.smembers(REDIS_SUBS_KEY)]
await send_post(post[0], tag_list)
@@ -345,7 +345,7 @@ async def send_callback(cq: CallbackQuery):
async def background_on_start():
await redis.delete('e621:update')
await redis.delete(REDIS_LOCK_KEY)
while True:
logging.warning('Checking updates...')
try: