jiemianNews/plugins/track.py

import contextlib
import traceback

from asyncio import sleep
from random import uniform

from pyrogram.enums import ParseMode
from pyrogram.errors import FloodWait
from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton

from ci import app, scheduler, channel_id, admin_id
from pyrogram import Client, filters

from defs.get_news import get_news
from defs.model import New


def gen_button(new: New):
    return InlineKeyboardMarkup(
        [
            [InlineKeyboardButton("🔗  链接", url=new.url)]
        ])


async def send_new_text(new: New):
    await app.send_message(channel_id, new.text, parse_mode=ParseMode.HTML, reply_markup=gen_button(new))


async def send_new_photo(new: New):
    if new.img_url:
        try:
            return await app.send_photo(channel_id, new.img_url, caption=new.text, parse_mode=ParseMode.HTML,
                                        reply_markup=gen_button(new))
        except Exception:
            return await send_new_text(new)
    return await send_new_text(new)


async def send_new(new: New):
    return await send_new_photo(new) if new.img_url else await send_new_text(new)


async def update_track(time: int = None):
    news = await get_news(time)
    news = sorted(news, key=lambda x: x.publish_time, reverse=False)
    for new in news:
        try:
            await send_new(new)
            print(f"{new.title} sent")
        except FloodWait as e:
            print(f"Send document flood - Sleep for {e.value} second(s)")
            await sleep(e.value + uniform(0.5, 1.0))
            await send_new(new)
        except Exception:
            traceback.print_exc()
            with contextlib.suppress(Exception):
                await app.send_message(admin_id, f"Error: {traceback.format_exc()}")
        await sleep(uniform(0.5, 2.0))
    print("run ok!")


@scheduler.scheduled_job("cron", minute="*/10", id="track")
async def run_every_10_minute():
    await update_track()


@Client.on_message(filters.incoming & filters.private & filters.chat(admin_id) &
                   filters.command(["force_update", ]))
async def force_update(_: Client, message: Message):
    try:
        a = int(message.text.replace("/force_update", "").strip())
    except ValueError:
        a = None
    await update_track(a)
    await message.reply("ok!")