diff --git a/defs/weather_graph.py b/defs/weather_graph.py new file mode 100644 index 0000000..3d0cd3f --- /dev/null +++ b/defs/weather_graph.py @@ -0,0 +1,70 @@ +import asyncio +import os +import shutil +from pathlib import Path + +import aiofiles +import httpx +from bs4 import BeautifulSoup + +from defs.bilibili_download import execute + +FFMPEG_PATH = "ffmpeg" +URL = "http://www.nmc.cn/publish/satellite/fy4b-visible.htm" +TEMP_PATH = Path("data") / "weather_tmp" +TEMP_PATH.mkdir(exist_ok=True, parents=True) +OUTPUT_PATH = TEMP_PATH / "output.mp4" + + +class WeatherGraphError(Exception): + def __init__(self, error: str): + self.error = error + + +async def download_img(client, index: int, url: str): + try: + res = await client.get(url) + if res.status_code != 200: + print(f"下载天气图片失败,错误代码 {res.status_code}") + async with aiofiles.open( + TEMP_PATH / f"img{str(index).zfill(3)}.jpg", "wb" + ) as f: + await f.write(res.content) + except Exception as e: + print(f"下载天气图片失败 {url} {e}") + + +async def get_images(): + async with httpx.AsyncClient() as client: + res = await client.get(URL) + if res.status_code != 200: + raise WeatherGraphError(f"获取图片数据失败,错误代码 {res.status_code}") + html = res.content + soup = BeautifulSoup(html, "lxml") + elements = soup.find_all(class_="time") + tasks = [] + for index, element in enumerate(reversed(elements)): + url = element["data-img"] + url = url.replace("medium/", "") + task = asyncio.create_task(download_img(client, index, url)) + tasks.append(task) + await asyncio.gather(*tasks) + + +async def get_video(): + frame_rate = 6 + ffmpeg_cmd = f'{FFMPEG_PATH} -r {frame_rate} -i "{TEMP_PATH}/img%03d.jpg" -vf scale=1920:1080 -c:v libx264 -pix_fmt yuv420p -y "{TEMP_PATH}/output.mp4"' + _, code = await execute(ffmpeg_cmd) + if code != 0: + raise WeatherGraphError("生成视频失败") + + +async def gen() -> Path: + if not os.path.exists(TEMP_PATH): + os.mkdir(TEMP_PATH) + else: + shutil.rmtree(TEMP_PATH) + os.mkdir(TEMP_PATH) + await get_images() + await get_video() + return OUTPUT_PATH diff --git a/modules/weather_graph.py b/modules/weather_graph.py new file mode 100644 index 0000000..5576cb4 --- /dev/null +++ b/modules/weather_graph.py @@ -0,0 +1,40 @@ +from pyrogram import filters, Client +from pyrogram.types import Message + +from init import bot, logger + +from cashews import cache +from defs.weather_graph import OUTPUT_PATH, gen, WeatherGraphError +from scheduler import scheduler + +CACHE_KEY = "weather_graph:file_id" + + +@bot.on_message( + filters.incoming + & filters.command(["weather_graph", f"weather_graph@{bot.me.username}"]) +) +async def weather_graph_command(_: "Client", message: "Message"): + if file_id := await cache.get(CACHE_KEY): + await message.reply_video(file_id, quote=True) + return + if OUTPUT_PATH.exists(): + reply = await message.reply_video(OUTPUT_PATH, quote=True) + await cache.set(CACHE_KEY, reply.video.file_id, expire=3600) + else: + reply = await message.reply("正在生成中...") + try: + await gen() + r1 = await message.reply_video(OUTPUT_PATH, quote=True) + await cache.set(CACHE_KEY, r1.video.file_id, expire=3600) + await reply.delete() + except WeatherGraphError as e: + await reply.edit(f"生成失败:{e.error}") + + +@scheduler.scheduled_job("interval", hours=1, id="weather_graph_refresh") +async def weather_graph_refresh(): + try: + await gen() + except WeatherGraphError as e: + logger.exception("Weather Graph Refresh failed", exc_info=e)