diff --git a/messages/base.py b/messages/base.py index f990078..623d310 100644 --- a/messages/base.py +++ b/messages/base.py @@ -13,6 +13,12 @@ class Base: def extra_info(self): return dict() + @property + def room_id(self): + if hasattr(self.instance, 'common'): + return self.instance.common.roomId + return None + def user(self): if(hasattr(self.instance, 'user')): return self.instance.user diff --git a/messages/chat.py b/messages/chat.py index 8d54b74..98aff9b 100644 --- a/messages/chat.py +++ b/messages/chat.py @@ -7,6 +7,10 @@ class ChatMessage(Base): def __init__(self): self.instance = message_pb2.ChatMessage() + @property + def content(self): + return self.instance.content + def format_content(self): return self.user().nickname + ': ' + self.instance.content diff --git a/messages/gift.py b/messages/gift.py index e96fd47..65b57c7 100644 --- a/messages/gift.py +++ b/messages/gift.py @@ -11,9 +11,14 @@ class GiftMessage(Base): return { 'giftId': self.instance.gift.id, 'giftName': self.instance.gift.name, - 'giftCount': self.instance.gift.diamondCount, + 'giftCount': self.instance.repeatCount, + 'diamondCount': self.instance.gift.diamondCount, } + @property + def gift(self): + return self.instance.gift + def format_content(self): return self.instance.common.describe diff --git a/output/DebugWriter.py b/output/DebugWriter.py deleted file mode 100644 index 11ff3c4..0000000 --- a/output/DebugWriter.py +++ /dev/null @@ -1,6 +0,0 @@ -from output import IOutput - - -class DebugWriter(IOutput): - def other_output(self, message_type: str, message_raw: bytes): - print(message_type) diff --git a/output/__init__.py b/output/__init__.py index 468ecb6..eed7200 100644 --- a/output/__init__.py +++ b/output/__init__.py @@ -1,6 +1,13 @@ -from output.IOutput import IOutput from output.print import Print +from output.debug import DebugWriter +from output.xml import XMLWriter +from typing import TYPE_CHECKING -OUTPUTER = [ - Print() -] \ No newline at end of file +if TYPE_CHECKING: + from output.IOutput import IOutput + +OUTPUTER: "list[IOutput]" = [ + Print(), + XMLWriter(), + DebugWriter() +] diff --git a/output/debug.py b/output/debug.py new file mode 100644 index 0000000..1960be2 --- /dev/null +++ b/output/debug.py @@ -0,0 +1,26 @@ +import os +import time +import traceback +from output.IOutput import IOutput + + +class DebugWriter(IOutput): + def other_output(self, message_type: str, message_raw: bytes): + if not os.path.isdir(os.path.join("", "debug")): + os.makedirs(os.path.join("", "debug")) + if not os.path.isdir(os.path.join("", "debug", message_type)): + os.makedirs(os.path.join("", "debug", message_type)) + with open(os.path.join("", "debug", message_type, str(time.time())), "wb") as f: + f.write(message_raw) + + def error_output(self, message_type: str, message_raw: bytes, exception: Exception): + if not os.path.isdir(os.path.join("", "error")): + os.makedirs(os.path.join("", "error")) + if not os.path.isdir(os.path.join("", "error", message_type)): + os.makedirs(os.path.join("", "error", message_type)) + ts = time.time() + with open(os.path.join("", "error", message_type, str(ts)), "wb") as f: + f.write(message_raw) + traceback.print_exc(file=open(os.path.join("", "error", message_type, str(ts)) + ".exc", "w", encoding="UTF-8")) + + diff --git a/output/print.py b/output/print.py index 23fcf55..3669a11 100644 --- a/output/print.py +++ b/output/print.py @@ -1,6 +1,6 @@ from colorama import init, Fore -from output import IOutput +from output.IOutput import IOutput RED = Fore.RED GREEN = Fore.GREEN diff --git a/output/xml.py b/output/xml.py new file mode 100644 index 0000000..eabb4fa --- /dev/null +++ b/output/xml.py @@ -0,0 +1,62 @@ +from output.IOutput import IOutput +from typing import IO +import time + + +class XMLWriter(IOutput): + """ + 可输出与B站弹幕姬兼容的xml弹幕格式,可用于转成ass字幕 + """ + def __init__(self): + self.file_mappings: "dict[str, IO[str]]" = {} + self.time_mappings: "dict[str, float]" = {} + + def _get_fd_by_room_id(self, room_id: str) -> IO[str]: + if room_id in self.file_mappings: + return self.file_mappings[room_id] + fd = open(f"{room_id}_{time.time()}.xml", "w", encoding="UTF-8") + self.file_mappings[room_id] = fd + self.time_mappings[room_id] = time.time() + return fd + + def _close_fd_by_room_id(self, room_id: str): + if room_id in self.file_mappings: + fd = self.file_mappings[room_id] + if not fd.closed: + fd.close() + del self.file_mappings[room_id] + if room_id in self.time_mappings: + del self.time_mappings[room_id] + + def control_output(self, message): + # 下播了 + self._close_fd_by_room_id(message.room_id) + + def _get_bias_ts_by_room_id(self, room_id: str, cur_ts: float = 0): + if cur_ts == 0: + cur_ts = time.time() + if room_id not in self.time_mappings: + return 0 + return cur_ts - self.time_mappings[room_id] + + def chat_output(self, message): + fd = self._get_fd_by_room_id(message.room_id) + if fd is None: + return + cur_time = time.time() + _c = """{}\r\n""".format( + self._get_bias_ts_by_room_id(message.room_id, cur_time), + cur_time*1000, message.user().id, message.user().nickname, message.content + ) + fd.write(_c) + + def gift_output(self, message): + fd = self._get_fd_by_room_id(message.room_id) + if fd is None: + return + cur_time = time.time() + _c = """\r\n""".format( + self._get_bias_ts_by_room_id(message.room_id, cur_time), + message.user().nickname, message.gift.name, message.instance.repeatCount + ) + fd.write(_c)