xml output

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

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
 

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
 

output/DebugWriter.py

@@ -1,6 +0,0 @@
-from output import IOutput
-
-
-class DebugWriter(IOutput):
-    def other_output(self, message_type: str, message_raw: bytes):
-        print(message_type)

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 = [
+if TYPE_CHECKING:
-    Print()
+    from output.IOutput import IOutput
-]
+
+OUTPUTER: "list[IOutput]" = [
+    Print(),
+    XMLWriter(),
+    DebugWriter()
+]

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"))
+
+

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

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 = """<d p="{:.2f},1,24,16777215,{:.0f},0,{},0" user="{}">{}</d>\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 = """<gift ts="{:.2f}" user="{}" giftname="{}" giftcount="{}"></gift>\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)