Commit cc9db51

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-05-28 06:48:59
fix(danmu): avoid OOM
1 parent eb5a8a6
Changed files (1)
src
others
src/others/danmu.py
@@ -11,11 +11,11 @@ from loguru import logger
 from pyrogram.client import Client
 from pyrogram.types import Message
 
-from config import DANMU, PREFIX, TZ
+from config import DANMU, PREFIX, TEXT_LENGTH, TZ
 from messages.parser import parse_msg
 from messages.progress import modify_progress
 from messages.sender import send2tg
-from messages.utils import blockquote, equal_prefix, smart_split, startswith_prefix
+from messages.utils import blockquote, equal_prefix, startswith_prefix
 from networking import hx_req
 from utils import number
 
@@ -102,22 +102,43 @@ async def query_danmu(client: Client, message: Message, *, full_history: bool =
     resp = await hx_req(DANMU.BASR_URL, "POST", data=payload, check_kv={"code": 0}, check_keys=["count", "data"], silent=True)
     count = resp["count"]
     if count == 0:
-        await modify_progress(message=status_msg, text=caption + "\n⚠️未匹配到任何弹幕, 请重新查询", force_update=True, **kwargs)
+        await modify_progress(message=status_msg, text=caption + "\n⚠️未找到匹配弹幕", force_update=True, **kwargs)
         return
-    all_danmu = resp["data"]
+    header = f"📖**弹幕记录 ({user})**:" if user else "📖**弹幕记录**:"
+    super_chats = defaultdict(Decimal)  # {"currency": amount}
+
+    parsed = parse_danmu(resp["data"], super_chats, show_name=show_name)
+    danmu = parsed["danmu"]
+    processed = parsed["num_messages"]
+    await modify_progress(message=status_msg, text=caption + f"\n⏳已获取 {processed} 条弹幕", force_update=True, **kwargs)
     if full_history:
-        while len(all_danmu) < count:
+        while processed < count:
             page += 1
             payload["page"] = page
             logger.debug(f"Query: {payload}")
             resp = await hx_req(DANMU.BASR_URL, "POST", data=payload, check_kv={"code": 0}, check_keys=["count", "data"], silent=True)
-            all_danmu.extend(resp["data"])
+            parsed = parse_danmu(resp["data"], super_chats, show_name=show_name)
+            danmu += parsed["danmu"]
+            processed += parsed["num_messages"]
+            await modify_progress(message=status_msg, text=caption + f"\n⏳已获取 {processed} 条弹幕", **kwargs)
 
-    header = f"📖**弹幕记录 ({user})**:" if user else "📖**弹幕记录**:"
+    profit = ""
+    for currency, amount in sorted(super_chats.items()):
+        profit += f"\n**{currency}**: {number(amount)}"
+    final = f"{header}{profit}\n{blockquote(danmu)}"
+    if not full_history or len(final) < TEXT_LENGTH:
+        await modify_progress(message=status_msg, text=blockquote(final), force_update=True, **kwargs)
+    else:
+        caption += f"\n#️⃣弹幕数: {count}"
+        caption += profit
+        with BytesIO(danmu.encode("utf-8")) as f:
+            await client.send_document(info["cid"], f, file_name=f"{user}弹幕记录.txt", caption=caption)
+
+
+def parse_danmu(data: list[dict], super_chats: defaultdict, *, show_name: bool = True) -> dict:
     msg = ""
-    super_chats = defaultdict(Decimal)  # {"currency": amount}
     sc_pattern = re.compile(r"^([A-Z]{2,}) (\d+(?:\.\d+)?)$")
-    for danmu in sorted(all_danmu, key=lambda x: x["timestamp"], reverse=True):  # new to old
+    for danmu in sorted(data, key=lambda x: x["timestamp"], reverse=True):  # new to old
         dt = datetime.fromtimestamp(danmu["timestamp"] / 1000000, tz=UTC).astimezone(ZoneInfo(TZ))
         sc_amount = ""
         if amt := danmu.get("scAmount"):
@@ -126,16 +147,4 @@ async def query_danmu(client: Client, message: Message, *, full_history: bool =
                 currency = matched.group(1)
                 super_chats[currency] += Decimal(matched.group(2))
         msg += f"\n{dt:%m-%d %H:%M:%S}|{danmu['authorName']}{sc_amount}: {danmu['message']}" if show_name else f"\n{dt:%m-%d %H:%M:%S}{sc_amount}: {danmu['message']}"
-    msg = msg.strip()
-    profit = ""
-    for currency, amount in sorted(super_chats.items()):
-        profit += f"\n**{currency}**: {number(amount)}"
-    final = f"{header}{profit}\n{blockquote(msg)}"
-    if not full_history or len(await smart_split(final)) == 1:
-        final = (await smart_split(final))[0]
-        await modify_progress(message=status_msg, text=blockquote(final), force_update=True, **kwargs)
-    else:
-        caption += f"\n#️⃣弹幕数: {count}"
-        caption += profit
-        with BytesIO(msg.encode("utf-8")) as f:
-            await client.send_document(info["cid"], f, file_name=f"{user}弹幕记录.txt", caption=caption)
+    return {"danmu": msg.strip(), "num_messages": len(data)}