initial commit

2023-03-22 18:07:42 +03:00
commit 4737cdee63
5 changed files with 1203 additions and 0 deletions
--- a/main.py
+++ b/main.py
@@ -0,0 +1,103 @@
+import asyncio
+
+import redis.asyncio as aioredis
+import logging
+from aiogram import Bot, Dispatcher
+from aiogram.types import Message
+from aiogram.utils import executor
+
+from e621 import E621
+
+redis = aioredis.from_url('redis://localhost')
+e621 = E621()
+
+logging.basicConfig(level=logging.INFO)
+bot = Bot(token='450146961:AAFqwPcXuWXaYP8j8sE_rYgp_dAsk9YWf7I')
+dp = Dispatcher(bot)
+
+
+async def check_updates():
+    async with redis.lock('e621:update'):
+        tag_list = [t.decode() for t in await redis.smembers('e621:subs')]
+        tags = ' '.join(f'~{tag}' for tag in tag_list)
+        posts = await e621.get_posts(tags)
+        if not posts:
+            return
+        already_sent = await redis.smismember('e621:sent', [p.id for p in posts])
+        for i in range(len(posts)):
+            if already_sent[i]:
+                continue
+            post = posts[i]
+            monitored_tags = set(post.tags.flatten()) & set(tag_list)
+            caption = f'Monitored tags: {" ".join(monitored_tags)}\n' \
+                      f'Artist: {" ".join(post.tags.artist)}\n' \
+                      f'Character: {" ".join(post.tags.character)}\n\n' \
+                      f'https://e621.net/posts/{post.id}'
+            if post.file.ext == 'webm':
+                await bot.send_video(98934915,
+                                     post.file.url,
+                                     width=post.file.width,
+                                     height=post.file.height,
+                                     thumb=post.preview.url,
+                                     caption=caption)
+            else:
+                await bot.send_photo(98934915,
+                                     post.file.url,
+                                     caption=caption)
+            await redis.sadd('e621:sent', post.id)
+
+
+@dp.message_handler(commands=['add'])
+async def add_tag(msg: Message):
+    args = msg.get_args()
+    if not args:
+        await msg.reply('Please provide tag to subscribe to')
+        return
+    for tag in args.split():
+        await redis.sadd('e621:subs', tag)
+    await msg.reply(f'Tags {args} added')
+
+
+@dp.message_handler(regexp=r'^\/del_\S+$')
+async def del_tag(msg: Message):
+    args = msg.text[5:]
+    if not args:
+        await msg.reply('Please provide tag to subscribe to')
+        return
+    if ' ' in args:
+        await msg.reply('Tag should not contain spaces')
+        return
+    if not await redis.sismember('e621:subs', args):
+        await msg.reply('Tag not found')
+        return
+    await redis.srem('e621:subs', args)
+    await msg.reply(f'Tag {msg} removed')
+
+
+@dp.message_handler(commands=['list'])
+async def list_tags(msg: Message):
+    tags = [t.decode() for t in await redis.smembers('e621:subs')]
+    lines = []
+    for tag in tags:
+        lines.append(f'- {tag} [/del_{tag}]')
+    lines = "\n".join(lines)
+    await msg.reply(f'Monitored tags:\n\n{lines}')
+
+
+@dp.message_handler(commands=['update'])
+async def update(msg: Message):
+    await check_updates()
+
+
+async def background_on_start():
+    while True:
+        await check_updates()
+        await asyncio.sleep(60)
+
+
+async def on_bot_startup(dp: Dispatcher):
+    asyncio.create_task(background_on_start())
+
+
+if __name__ == '__main__':
+    executor.start_polling(dp, on_startup=on_bot_startup)