diff --git a/src/blrec/core/danmaku_dumper.py b/src/blrec/core/danmaku_dumper.py index 17f0592..d158e94 100644 --- a/src/blrec/core/danmaku_dumper.py +++ b/src/blrec/core/danmaku_dumper.py @@ -9,7 +9,7 @@ from tenacity import AsyncRetrying, retry_if_not_exception_type, stop_after_atte from blrec import __github__, __prog__, __version__ from blrec.bili.live import Live -from blrec.core.models import GiftSendMsg, GuardBuyMsg, SuperChatMsg +from blrec.core.models import GiftSendMsg, GuardBuyMsg, SuperChatMsg, UserToastMsg from blrec.danmaku.io import DanmakuWriter from blrec.danmaku.models import ( Danmu, @@ -17,6 +17,7 @@ from blrec.danmaku.models import ( GuardBuyRecord, Metadata, SuperChatRecord, + UserToast, ) from blrec.event.event_emitter import EventEmitter, EventListener from blrec.exception import exception_callback, submit_exception @@ -198,6 +199,8 @@ class DanmakuDumper( if isinstance(msg, DanmuMsg): await writer.write_danmu(self._make_danmu(msg)) self._statistics.submit(1) + elif isinstance(msg, UserToastMsg): + await writer.write_user_toast(self._make_user_toast(msg)) elif isinstance(msg, GiftSendMsg): if not self.record_gift_send: continue @@ -288,5 +291,18 @@ class DanmakuDumper( message=msg.message, ) + def _make_user_toast(self, msg: UserToastMsg) -> UserToast: + return UserToast( + ts=self._calc_stime(msg.start_time * 1000), + uid=msg.uid, + user=msg.username, + unit=msg.unit, + count=msg.num, + price=msg.price, + role=msg.role_name, + level=msg.guard_level, + msg=msg.toast_msg, + ) + def _calc_stime(self, timestamp: int) -> float: return (max(timestamp - self._timebase, 0) + self._delta) / 1000 diff --git a/src/blrec/core/danmaku_receiver.py b/src/blrec/core/danmaku_receiver.py index 3aabc4c..d96ece5 100644 --- a/src/blrec/core/danmaku_receiver.py +++ b/src/blrec/core/danmaku_receiver.py @@ -6,7 +6,7 @@ from blrec.bili.danmaku_client import DanmakuClient, DanmakuCommand, DanmakuList from blrec.bili.typing import Danmaku from blrec.utils.mixins import StoppableMixin -from .models import DanmuMsg, GiftSendMsg, GuardBuyMsg, SuperChatMsg +from .models import DanmuMsg, GiftSendMsg, GuardBuyMsg, SuperChatMsg, UserToastMsg from .typing import DanmakuMsg __all__ = ('DanmakuReceiver',) @@ -47,6 +47,8 @@ class DanmakuReceiver(DanmakuListener, StoppableMixin): msg = GuardBuyMsg.from_danmu(danmu) elif cmd == DanmakuCommand.SUPER_CHAT_MESSAGE.value: msg = SuperChatMsg.from_danmu(danmu) + elif cmd == DanmakuCommand.USER_TOAST_MSG.value: + msg = UserToastMsg.from_danmu(danmu) else: return diff --git a/src/blrec/core/models.py b/src/blrec/core/models.py index f5e1323..4d8b8f6 100644 --- a/src/blrec/core/models.py +++ b/src/blrec/core/models.py @@ -38,6 +38,34 @@ class DanmuMsg: ) +@attr.s(auto_attribs=True, slots=True, frozen=True) +class UserToastMsg: + start_time: int # timestamp in seconds + uid: int + username: str + unit: str + num: int + price: int + role_name: str + guard_level: str + toast_msg: str + + @staticmethod + def from_danmu(danmu: Danmaku) -> 'UserToastMsg': + data = danmu['data'] + return UserToastMsg( + start_time=data['start_time'], + uid=data['uid'], + username=data['username'], + unit=data['unit'], + num=data['num'], + price=data['price'], + role_name=data['role_name'], + guard_level=data['guard_level'], + toast_msg=data['toast_msg'].replace('<%', '').replace('%>', ''), + ) + + @attr.s(auto_attribs=True, frozen=True, slots=True) class GiftSendMsg: gift_name: str diff --git a/src/blrec/core/typing.py b/src/blrec/core/typing.py index b8a8951..be1c005 100644 --- a/src/blrec/core/typing.py +++ b/src/blrec/core/typing.py @@ -1,12 +1,5 @@ from typing import Union +from .models import DanmuMsg, GiftSendMsg, GuardBuyMsg, SuperChatMsg, UserToastMsg -from .models import DanmuMsg, GiftSendMsg, GuardBuyMsg, SuperChatMsg - - -DanmakuMsg = Union[ - DanmuMsg, - GiftSendMsg, - GuardBuyMsg, - SuperChatMsg, -] +DanmakuMsg = Union[DanmuMsg, GiftSendMsg, GuardBuyMsg, SuperChatMsg, UserToastMsg] diff --git a/src/blrec/danmaku/io.py b/src/blrec/danmaku/io.py index 2b1d390..c63c68d 100644 --- a/src/blrec/danmaku/io.py +++ b/src/blrec/danmaku/io.py @@ -11,7 +11,14 @@ import aiofiles import attr from lxml import etree -from .models import Danmu, GiftSendRecord, GuardBuyRecord, Metadata, SuperChatRecord +from .models import ( + Danmu, + GiftSendRecord, + GuardBuyRecord, + Metadata, + SuperChatRecord, + UserToast, +) from .typing import Element __all__ = 'DanmakuReader', 'DanmakuWriter' @@ -119,6 +126,9 @@ class DanmakuWriter: async def write_danmu(self, danmu: Danmu) -> None: await self._file.write(self._serialize_danmu(danmu)) + async def write_user_toast(self, toast: UserToast) -> None: + await self._file.write(self._serialize_user_toast(toast)) + async def write_gift_send_record(self, record: GiftSendRecord) -> None: await self._file.write(self._serialize_gift_send_record(record)) @@ -178,6 +188,16 @@ class DanmakuWriter: return ' ' + etree.tostring(elem, encoding='utf8').decode() + '\n' + def _serialize_user_toast(self, toast: UserToast) -> str: + attrib = attr.asdict( + toast, + filter=lambda a, v: a.name != 'msg', + value_serializer=record_value_serializer, + ) + elem = etree.Element('toast', attrib=attrib) + elem.text = remove_control_characters(toast.msg) + return ' ' + etree.tostring(elem, encoding='utf8').decode() + '\n' + def _serialize_gift_send_record(self, record: GiftSendRecord) -> str: attrib = attr.asdict(record, value_serializer=record_value_serializer) elem = etree.Element('gift', attrib=attrib) @@ -195,7 +215,7 @@ class DanmakuWriter: value_serializer=record_value_serializer, ) elem = etree.Element('sc', attrib=attrib) - elem.text = record.message + elem.text = remove_control_characters(record.message) return ' ' + etree.tostring(elem, encoding='utf8').decode() + '\n' diff --git a/src/blrec/danmaku/models.py b/src/blrec/danmaku/models.py index dad7256..783d3dc 100644 --- a/src/blrec/danmaku/models.py +++ b/src/blrec/danmaku/models.py @@ -30,6 +30,19 @@ class Danmu: text: str +@attr.s(auto_attribs=True, slots=True, frozen=True) +class UserToast: + ts: float # start_time + uid: int + user: str # username + unit: str + count: int # num + price: int + role: str # role_name + level: str # guard_level + msg: str # toast_msg + + @attr.s(auto_attribs=True, slots=True, frozen=True) class GiftSendRecord: ts: float