✨ 支持 download cache

2024-11-21 14:48:23 +00:00 · 2023-08-18 22:13:02 +08:00 · 2023-08-18 22:13:02 +08:00 · b40ee45919
commit b40ee45919
parent 6a23509ae3
6 changed files with 138 additions and 24 deletions
--- a/defs/bilibili_download.py
+++ b/defs/bilibili_download.py
@ -15,6 +15,8 @@ from pyrogram.types import Message

 from defs.request import cache_dir
 from init import bot, logger, request
+from models.models.bilifav import BiliFav
+from models.services.bilifav import BiliFavAction

 FFMPEG_PATH = "ffmpeg"
 FFPROBE_PATH = "ffprobe"
@ -217,6 +219,19 @@ async def take_screenshot(info: Dict) -> Optional[BytesIO]:
        return None


+def gen_audio_caption(a: Audio, info: Dict) -> str:
+    intro = info.get("intro", "")
+    if intro:
+        text = f"<b>{info['title']}</b>\n\n{intro}\n\nhttps://www.bilibili.com/audio/au{a.get_auid()}"
+        if len(text) > 800:
+            text = f"<b>{info['title']}</b>\n\n简介过长，无法显示\n\nhttps://www.bilibili.com/audio/au{a.get_auid()}"
+    else:
+        text = (
+            f"<b>{info['title']}</b>\n\nhttps://www.bilibili.com/audio/au{a.get_auid()}"
+        )
+    return text
+
+
 async def audio_download(
    a: Audio, m: Message, push_id: int = None
 ) -> Optional[Message]:
@ -238,17 +253,11 @@ async def audio_download(
                thumb.seek(0)
            else:
                thumb = None
-        intro = info.get("intro", "")
-        if intro:
-            text = f"<b>{info['title']}</b>\n\n{intro}\n\nhttps://www.bilibili.com/audio/au{a.get_auid()}"
-            if len(text) > 800:
-                text = f"<b>{info['title']}</b>\n\n简介过长，无法显示\n\nhttps://www.bilibili.com/audio/au{a.get_auid()}"
-        else:
-            text = f"<b>{info['title']}</b>\n\nhttps://www.bilibili.com/audio/au{a.get_auid()}"
+        caption = gen_audio_caption(a, info)
        msg = await bot.send_audio(
            chat_id=push_id or m.chat.id,
            audio=media,
-            caption=text,
+            caption=caption,
            parse_mode=ParseMode.HTML,
            reply_to_message_id=m.reply_to_message_id if not push_id else None,
            thumb=thumb,
@ -256,6 +265,18 @@ async def audio_download(
            duration=info.get("duration"),
            performer=info.get("author"),
        )
+        if info.get("id") and (not await BiliFavAction.get_by_id(info.get("id"))):
+            audio_db = BiliFav(
+                id=info.get("id", 0),
+                bv_id=info.get("bvid").lower(),
+                type=12,
+                title=info.get("title", ""),
+                cover=info.get("cover", ""),
+                message_id=0,
+                file_id=msg.audio.file_id,
+                timestamp=int(time.time()),
+            )
+            await BiliFavAction.add_bili_fav(audio_db)
    except BilibiliDownloaderError as e:
        await fail_edit(m, e.MSG)
        return
@ -336,6 +357,17 @@ async def go_upload_progress(current: int, total: int, m: Message):
        await message_edit(total, current, chunk, chunk_time, m, "上传")


+def gen_video_caption(v: Video, info: Dict) -> str:
+    caption = (
+        f"<b>{info['title']}</b>\n\n{info['desc']}\n\nhttps://b23.tv/{v.get_bvid()}"
+    )
+    if len(caption) > 800:
+        caption = (
+            f"<b>{info['title']}</b>\n\n简介过长，无法显示\n\nhttps://b23.tv/{v.get_bvid()}"
+        )
+    return caption
+
+
 async def go_upload(
    v: Video, p_num: int, m: Message, push_id: int = None
 ) -> Optional[Message]:
@ -349,10 +381,9 @@ async def go_upload(
        try:
            info = await v.get_info()
            video_jpg = await take_screenshot(info)
-            caption = f"<b>{info['title']}</b>\n\n{info['desc']}\n\nhttps://b23.tv/{v.get_bvid()}"
-            if len(caption) > 800:
-                caption = f"<b>{info['title']}</b>\n\n简介过长，无法显示\n\nhttps://b23.tv/{v.get_bvid()}"
+            caption = gen_video_caption(v, info)
        except Exception:
+            info = None
            video_jpg = None
            caption = f"https://b23.tv/{v.get_bvid()}"
        logger.info(f"Uploading {video_path}")
@ -370,6 +401,22 @@ async def go_upload(
            progress_args=(m,),
            reply_to_message_id=m.reply_to_message_id if not push_id else None,
        )
+        if (
+            (not await BiliFavAction.get_by_bv_id(v.get_bvid()))
+            and info is not None
+            and info.get("aid")
+        ):
+            video_db = BiliFav(
+                id=info.get("aid", 0),
+                bv_id=info.get("bvid").lower(),
+                type=2,
+                title=info.get("title", ""),
+                cover=info.get("pic", ""),
+                message_id=0,
+                file_id=msg.video.file_id,
+                timestamp=int(time.time()),
+            )
+            await BiliFavAction.add_bili_fav(video_db)
        logger.info(f"Upload {video_path} success")
    except BilibiliDownloaderError as e:
        await fail_edit(m, e.MSG)
--- a/defs/bilibili_fav.py
+++ b/defs/bilibili_fav.py
@ -2,16 +2,24 @@ import time
 from enum import Enum
 from typing import Optional

+from bilibili_api.audio import Audio
 from bilibili_api.favorite_list import FavoriteList
+from bilibili_api.video import Video
 from pydantic import BaseModel, ValidationError
 from pyrogram.types import Message

 from defs.bilibili import credential, create_video, create_audio
-from defs.bilibili_download import go_download, go_upload, audio_download
+from defs.bilibili_download import (
+    go_download,
+    go_upload,
+    audio_download,
+    gen_video_caption,
+    gen_audio_caption,
+)
 from defs.glover import bilifav_id, bilifav_channel
 from models.models.bilifav import BiliFav
 from models.services.bilifav import BiliFavAction
-from init import logger
+from init import logger, bot

 fav = FavoriteList(media_id=bilifav_id, credential=credential)

@ -36,9 +44,25 @@ class Media(BaseModel):
    """ 简介 """


+async def process_video_from_cache(video: Video, video_db: BiliFav):
+    info = await video.get_info()
+    caption = gen_video_caption(video, info)
+    msg = await bot.send_video(
+        bilifav_channel,
+        video_db.file_id,
+        caption=caption,
+    )
+    video_db.message_id = msg.id
+    video_db.timestamp = int(time.time())
+    await BiliFavAction.update_bili_fav(video_db)
+
+
 async def process_video(data: Media, m: Message):
    """处理视频"""
    video = create_video(data.bvid)
+    if video_db := await BiliFavAction.get_by_bv_id(data.bvid):
+        await process_video_from_cache(video, video_db)
+        return
    await go_download(video, 0, m, task=False)
    msg = await go_upload(video, 0, m, push_id=bilifav_channel)
    if not msg:
@ -56,9 +80,25 @@ async def process_video(data: Media, m: Message):
    await BiliFavAction.add_bili_fav(video_db)


+async def process_audio_from_cache(audio: Audio, audio_db: BiliFav):
+    info = await audio.get_info()
+    caption = gen_audio_caption(audio, info)
+    msg = await bot.send_audio(
+        bilifav_channel,
+        audio_db.file_id,
+        caption=caption,
+    )
+    audio_db.message_id = msg.id
+    audio_db.timestamp = int(time.time())
+    await BiliFavAction.update_bili_fav(audio_db)
+
+
 async def process_audio(data: Media, m: Message):
    """处理音频"""
    audio = create_audio(f"au{data.id}")
+    if audio_db := await BiliFavAction.get_by_bv_id(data.bvid):
+        await process_audio_from_cache(audio, audio_db)
+        return
    msg = await audio_download(audio, m, push_id=bilifav_channel)
    if not msg:
        raise BilibiliFavException
@ -90,7 +130,7 @@ async def check_update(m: Message):
        except ValidationError as _:
            logger.exception("Validate media failed")
            continue
-        if await BiliFavAction.get_by_bv_id(data.bvid):
+        if await BiliFavAction.get_by_bv_id(data.bvid, fav=True):
            continue
        n = await m.reply(f"处理 {data.type.name} {data.bvid} 中...")
        try:
--- a/models/models/bilifav.py
+++ b/models/models/bilifav.py
@ -1,3 +1,5 @@
+from typing import Optional
+
 from sqlmodel import SQLModel, Field


@ -9,6 +11,6 @@ class BiliFav(SQLModel, table=True):
    type: int = Field(default=2)
    title: str = Field()
    cover: str = Field()
-    message_id: int = Field(default=0)
+    message_id: Optional[int] = Field(default=0, nullable=True)
    file_id: str = Field()
    timestamp: int = Field(default=0)
--- a/models/services/bilifav.py
+++ b/models/services/bilifav.py
@ -9,20 +9,24 @@ from models.models.bilifav import BiliFav

 class BiliFavAction:
    @staticmethod
-    async def get_by_id(id_: int) -> Optional[BiliFav]:
+    async def get_by_id(id_: int, fav: bool = False) -> Optional[BiliFav]:
        async with sqlite.session() as session:
            session = cast(AsyncSession, session)
            statement = select(BiliFav).where(BiliFav.id == id_)
+            if fav:
+                statement = statement.where(BiliFav.message_id != 0)
            results = await session.exec(statement)
            return post[0] if (post := results.first()) else None

    @staticmethod
-    async def get_by_bv_id(bv_id: str) -> Optional[BiliFav]:
+    async def get_by_bv_id(bv_id: str, fav: bool = False) -> Optional[BiliFav]:
        if not bv_id:
            return None
        async with sqlite.session() as session:
            session = cast(AsyncSession, session)
            statement = select(BiliFav).where(BiliFav.bv_id == bv_id.lower())
+            if fav:
+                statement = statement.where(BiliFav.message_id != 0)
            results = await session.exec(statement)
            return post[0] if (post := results.first()) else None

--- a/modules/bilibili_download.py
+++ b/modules/bilibili_download.py
@ -33,7 +33,9 @@ async def bili_download_resolve(_: Client, message: Message):
    if video_db := await BiliFavAction.get_by_bv_id(video.get_bvid()):
        await message.reply_video(
            video_db.file_id,
-            caption=f"详细信息：https://t.me/{bilifav_channel_username}/{video_db.message_id}",
+            caption=f"详细信息：https://t.me/{bilifav_channel_username}/{video_db.message_id}"
+            if video_db.message_id
+            else None,
            quote=True,
        )
        raise ContinuePropagation
@ -50,6 +52,15 @@ async def bili_audio_download_resolve(_: Client, message: Message):
    else:
        raise ContinuePropagation
    audio = create_audio(audio_number)
+    if audio_db := await BiliFavAction.get_by_id(audio.get_auid()):
+        await message.reply_audio(
+            audio_db.file_id,
+            caption=f"详细信息：https://t.me/{bilifav_channel_username}/{audio_db.message_id}"
+            if audio_db.message_id
+            else None,
+            quote=True,
+        )
+        raise ContinuePropagation
    m = await message.reply("开始获取音频数据", quote=True)
    bot.loop.create_task(audio_download(audio, m))

--- a/modules/bilibili_fav.py
+++ b/modules/bilibili_fav.py
@ -2,11 +2,15 @@ import re
 import time

 from pyrogram import filters, Client, ContinuePropagation
-from pyrogram.types import Message, CallbackQuery
+from pyrogram.types import Message

 from defs.bilibili import b23_extract, create_video, create_audio
 from defs.bilibili_download import go_download, go_upload, audio_download
-from defs.bilibili_fav import check_update
+from defs.bilibili_fav import (
+    check_update,
+    process_video_from_cache,
+    process_audio_from_cache,
+)
 from defs.glover import admin, bilifav_channel
 from init import bot, logger
 from models.models.bilifav import BiliFav
@ -15,10 +19,13 @@ from models.services.bilifav import BiliFavAction

 async def process_audio(video_number: str, message: Message):
    id_ = int(video_number[2:])
-    if await BiliFavAction.get_by_id(id_):
+    if await BiliFavAction.get_by_id(id_, fav=True):
        await message.reply("该音频已经存在")
        raise ContinuePropagation
    audio = create_audio(video_number)
+    if audio_db := await BiliFavAction.get_by_id(id_):
+        await process_audio_from_cache(audio, audio_db)
+        return
    info = await audio.get_info()
    m = await message.reply("开始获取音频数据", quote=True)
    msg = await audio_download(audio, m, push_id=bilifav_channel)
@ -39,9 +46,12 @@ async def process_audio(video_number: str, message: Message):

 async def process_video(video_number: str, p_num: int, message: Message):
    video = create_video(video_number)
-    if await BiliFavAction.get_by_bv_id(video.get_bvid()):
+    if await BiliFavAction.get_by_bv_id(video.get_bvid(), fav=True):
        await message.edit("该视频已经存在")
        raise ContinuePropagation
+    if video_db := await BiliFavAction.get_by_bv_id(video.get_bvid()):
+        await process_video_from_cache(video, video_db)
+        return
    info = await video.get_info()
    id_ = info.get("aid", 0)
    if not id_:
@ -52,7 +62,7 @@ async def process_video(video_number: str, p_num: int, message: Message):
    msg = await go_upload(video, p_num, m, push_id=bilifav_channel)
    if not msg:
        raise ContinuePropagation
-    audio_db = BiliFav(
+    video_db = BiliFav(
        id=id_,
        bv_id=info.get("bvid", "").lower(),
        type=2,
@ -62,7 +72,7 @@ async def process_video(video_number: str, p_num: int, message: Message):
        file_id=msg.video.file_id,
        timestamp=int(time.time()),
    )
-    await BiliFavAction.add_bili_fav(audio_db)
+    await BiliFavAction.add_bili_fav(video_db)


@bot.on_message(