From 788c8cba9d31dd6538f3ff57b802f6d9b989da1f Mon Sep 17 00:00:00 2001
From: Jerry Yan <792602257@qq.com>
Date: Sun, 5 Jun 2022 22:39:19 +0800
Subject: [PATCH] =?UTF-8?q?=E7=95=99=E4=B8=80=E4=B8=AA=E9=80=80=E5=87=BA?=
 =?UTF-8?q?=E5=90=8E=E6=B8=85=E7=90=86=E7=9A=84=E9=80=BB=E8=BE=91=EF=BC=8C?=
 =?UTF-8?q?=E8=99=BD=E7=84=B6windows=E4=B8=8A=E7=94=A8=E4=B8=8D=E4=BA=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.py           | 10 ++++++++++
 output/manager.py |  6 +++++-
 output/xml.py     | 27 +++++++++++++++++----------
 3 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/main.py b/main.py
index ca0edea..768b348 100644
--- a/main.py
+++ b/main.py
@@ -1,5 +1,6 @@
 import threading
 import subprocess
+import atexit
 
 from config.helper import config
 from handler.http_server import app
@@ -15,5 +16,14 @@ if __name__ == '__main__':
     api_thread.start()
     browser_manager = init_browser_manager()
     output_manager = OutputManager()
+
+
+    def terminate():
+        print("terminate")
+        browser_manager.terminate()
+        output_manager.terminate()
+
+
+    atexit.register(terminate)
     output_manager.start_loop()
     api_thread.join()
diff --git a/output/manager.py b/output/manager.py
index e241a91..47f637b 100644
--- a/output/manager.py
+++ b/output/manager.py
@@ -40,7 +40,7 @@ class OutputManager():
             self._writer.append(self._mapping[_c]())
 
     def __del__(self):
-        ...
+        self.terminate()
 
     def decode_payload(self, message: MessagePayload):
         try:
@@ -107,3 +107,7 @@ class OutputManager():
         while True:
             message = MESSAGE_QUEUE.get()
             self.decode_payload(message)
+
+    def terminate(self):
+        for writer in self._writer:
+            writer.terminate()
diff --git a/output/xml.py b/output/xml.py
index 885bfc6..ecb5804 100644
--- a/output/xml.py
+++ b/output/xml.py
@@ -8,29 +8,30 @@ class XMLWriter(IOutput):
     """
     可输出与B站弹幕姬兼容的xml弹幕格式，可用于转成ass字幕
     """
+
     def __init__(self):
-        self.file_mappings: "dict[str, IO[str]]" = {}
+        self._file_mappings: "dict[str, IO[str]]" = {}
         self.time_mappings: "dict[str, float]" = {}
         self._file_name_pattern: "str" = config()['output']['xml']['file_pattern']
 
     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]
+        if room_id in self._file_mappings:
+            return self._file_mappings[room_id]
         cur_ts = time.time()
         fd = open(self._file_name_pattern.format_map({
             "room_id": room_id,
             "ts": cur_ts
         }), "w", encoding="UTF-8")
-        self.file_mappings[room_id] = fd
+        self._file_mappings[room_id] = fd
         self.time_mappings[room_id] = cur_ts
         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 room_id in self._file_mappings:
+            fd = self._file_mappings[room_id]
             if not fd.closed:
                 fd.close()
-            del self.file_mappings[room_id]
+            del self._file_mappings[room_id]
         if room_id in self.time_mappings:
             del self.time_mappings[room_id]
 
@@ -50,9 +51,9 @@ class XMLWriter(IOutput):
         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
+        _c = """<d p="{:.2f},1,24,{},{:.0f},0,{},0" user="{}">{}</d>\r\n""".format(
+            self._get_bias_ts_by_room_id(message.room_id, cur_time), message.room_id,
+            cur_time * 1000, message.user().id, message.user().nickname, message.content
         )
         fd.write(_c)
 
@@ -66,3 +67,9 @@ class XMLWriter(IOutput):
             message.user().nickname, message.gift.name, message.instance.repeatCount
         )
         fd.write(_c)
+
+    def terminate(self):
+        print("保存所有弹幕文件中...")
+        for _room_id in self._file_mappings:
+            self._close_fd_by_room_id(_room_id)
+        print("保存完毕")