From 427f28bd5744f2cd2d22bd663e8df03c2e8f3025 Mon Sep 17 00:00:00 2001
From: Jerry Yan <792602257@qq.com>
Date: Sun, 5 Jun 2022 09:35:11 +0800
Subject: [PATCH] xml output
---
messages/base.py | 6 +++++
messages/chat.py | 4 +++
messages/gift.py | 7 ++++-
output/DebugWriter.py | 6 -----
output/__init__.py | 15 ++++++++---
output/debug.py | 26 ++++++++++++++++++
output/print.py | 2 +-
output/xml.py | 62 +++++++++++++++++++++++++++++++++++++++++++
8 files changed, 116 insertions(+), 12 deletions(-)
delete mode 100644 output/DebugWriter.py
create mode 100644 output/debug.py
create mode 100644 output/xml.py
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)