Reduce unnecessary levels

biliup/engine/download.py

@@ -1,16 +1,16 @@
+import asyncio
 import logging
 import os
 import re
 import subprocess
 import sys
+import threading
 import time
+from urllib.parse import urlparse
 
 import stream_gears
 
 from biliup.config import config
-from urllib.parse import urlparse
-
-from biliup.plugins.Danmaku.danmaku_main import Danmaku
 
 logger = logging.getLogger('biliup')
 
@@ -63,11 +63,13 @@ class DownloadBase:
         else:
             filename = f'{self.fname}%Y-%m-%dT%H_%M_%S'
         filename = get_valid_filename(filename)
+        fmtname = time.strftime(filename.encode("unicode-escape").decode()).encode().decode("unicode-escape")
+        threading.Thread(target=asyncio.run, args=(self.danmaku_download_start(fmtname),)).start()
         if self.downloader == 'stream-gears':
             stream_gears_download(self.raw_stream_url, self.fake_headers, filename, config.get('segment_time'),
                                   config.get('file_size'))
         else:
-            self.ffmpeg_download(filename)
+            self.ffmpeg_download(fmtname)
 
     def ffmpeg_download(self, filename):
         default_input_args = ['-headers', ''.join('%s: %s\r\n' % x for x in self.fake_headers.items()),
@@ -81,13 +83,12 @@ class DownloadBase:
                 '-c', 'copy', '-f', self.suffix]
         if config.get('segment_time'):
             args += ['-f', 'segment',
-                     f'{(time.strftime(filename).encode("unicode-escape").decode()).encode().decode("unicode-escape")} part-%03d.{self.suffix}']
+                     f'{filename} part-%03d.{self.suffix}']
         else:
             args += [
-                f'{(time.strftime(filename).encode("unicode-escape").decode()).encode().decode("unicode-escape")}.{self.suffix}.part']
+                f'{filename}.{self.suffix}.part']
 
         proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-        self.danmaku_download_start(filename)
         try:
             with proc.stdout as stdout:
                 for line in iter(stdout.readline, b''):  # b'\n'-separated lines
@@ -95,18 +96,13 @@ class DownloadBase:
                     print(decode_line, end='', file=sys.stderr)
                     logger.debug(decode_line.rstrip())
             retval = proc.wait()
-            self.danmaku_download_stop()
         except KeyboardInterrupt:
-            self.danmaku_download_stop()
             if sys.platform != 'win32':
                 proc.communicate(b'q')
             raise
         return retval
 
-    def danmaku_download_start(self, filename):
-        pass
-
-    def danmaku_download_stop(self):
+    async def danmaku_download_start(self, filename):
         pass
 
     def run(self):

biliup/plugins/Danmaku/__init__.py

@@ -1 +1,137 @@
+# 部分弹幕功能代码来自项目：https://github.com/IsoaSFlus/danmaku，感谢大佬
+# 快手弹幕代码来源及思路：https://github.com/py-wuhao/ks_barrage，感谢大佬
+# 仅抓取用户弹幕，不包括入场提醒、礼物赠送等。
 
+import asyncio
+import os
+import re
+import ssl
+import time
+
+import lxml.etree as etree
+import aiofiles
+import aiohttp
+
+from biliup.plugins.Danmaku.douyu import Douyu
+from biliup.plugins.Danmaku.huya import Huya
+from biliup.plugins.Danmaku.bilibili import Bilibili
+
+__all__ = ['DanmakuClient']
+
+
+class DanmakuClient:
+    def __init__(self, url, filename):
+        self.__starttime = time.time()
+        self.__filename = filename
+        self.__url = ''
+        self.__site = None
+        self.__hs = None
+        self.__ws = None
+        self.__stop = False
+        self.__dm_queue = None
+        self.__link_status = True
+
+        if 'http://' == url[:7] or 'https://' == url[:8]:
+            self.__url = url
+        else:
+            self.__url = 'http://' + url
+        for u, s in {'douyu.com': Douyu, 'huya.com': Huya, 'live.bilibili.com': Bilibili}.items():
+            if re.match(r'^(?:http[s]?://)?.*?%s/(.+?)$' % u, url):
+                self.__site = s
+                self.__u = u
+                break
+        if self.__site is None:
+            print('Invalid link!')
+            exit()
+
+    async def init_ws(self):
+        ws_url, reg_datas = await self.__site.get_ws_info(self.__url)
+        ctx = ssl.create_default_context()
+        ctx.set_ciphers('DEFAULT')
+        self.__ws = await self.__hs.ws_connect(ws_url, ssl_context=ctx)
+        if reg_datas:
+            for reg_data in reg_datas:
+                await self.__ws.send_bytes(reg_data)
+        if not os.path.exists(self.__filename):
+            async with aiofiles.open(self.__filename, mode='w') as f:
+                await f.write("<?xml version='1.0' encoding='UTF-8'?>\n"
+                              "<i xmlns:ns0='http://www.w3.org/1999/XSL/Transform'>\n"
+                              "</i>"
+                              )
+
+    async def heartbeats(self):
+        while not self.__stop and self.__site.heartbeat:
+            await asyncio.sleep(self.__site.heartbeatInterval)
+            try:
+                await self.__ws.send_bytes(self.__site.heartbeat)
+            except Exception as Error:
+                print(f"捕获到异常：{Error}")
+        await self.__hs.close()
+
+    async def fetch_danmaku(self):
+        # while not self.__stop:
+            await asyncio.sleep(1)
+            await self.init_ws()
+            await asyncio.sleep(1)
+            async for msg in self.__ws:
+                if self.__stop:
+                    await self.__dm_queue.put(None)
+                    return
+                self.__link_status = True
+                ms = self.__site.decode_msg(msg.data)
+                for m in ms:
+                    await self.__dm_queue.put(m)
+
+    async def print_danmaku(self):
+        parser = etree.XMLParser(recover=True)
+        tree = etree.parse(self.__filename, parser=parser)
+        root = tree.getroot()
+        msg_i = 0
+        msg_col = {'0': '16777215', '1': '16717077', '2': '2000880', '3': '8046667', '4': '16744192', '5': '10172916',
+                   '6': '16738740'}
+        while not self.__stop:
+            try:
+                m = await self.__dm_queue.get()
+                if m is None:
+                    return
+                # print(m)
+                if m['msg_type'] == 'danmaku':
+                    d = etree.SubElement(root, 'd')
+                    if 'col' in m:
+                        color = msg_col[m["col"]]
+                    elif 'color' in m:
+                        color = m["color"]
+                    # else:
+                    #     color = '16777215'
+                    msg_time = format(time.time() - self.__starttime, '.3f')
+                    d.set('p', f"{msg_time},1,25,{color},0,0,0,0")
+                    d.text = m["content"]
+
+                    if msg_i >= 5:
+                        etree.indent(root, "\t")
+                        tree.write(self.__filename, encoding="UTF-8", xml_declaration=True, pretty_print=True)
+                        msg_i = 0
+                    else:
+                        msg_i = msg_i + 1
+
+            except Exception as Error:
+                print(f"捕获到异常：{Error}")
+
+    async def start(self):
+        self.__dm_queue = asyncio.Queue()
+        self.__hs = aiohttp.ClientSession()
+        await self.init_ws()
+        await asyncio.gather(
+            self.heartbeats(),
+            self.fetch_danmaku(),
+            self.print_danmaku(),
+        )
+        print("弹幕录制结束")
+
+    def stop(self):
+        self.__stop = True
+
+
+
+# 虎牙直播：https://www.huya.com/lpl
+# 斗鱼直播：https://www.douyu.com/9999

biliup/plugins/Danmaku/danmaku_main.py

@@ -1,53 +0,0 @@
-# 部分弹幕功能代码来自项目：https://github.com/IsoaSFlus/danmaku，感谢大佬
-# 快手弹幕代码来源及思路：https://github.com/py-wuhao/ks_barrage，感谢大佬
-# 仅抓取用户弹幕，不包括入场提醒、礼物赠送等。
-
-import asyncio
-import queue
-import threading
-import time
-from concurrent.futures import ThreadPoolExecutor
-
-from biliup.plugins.Danmaku import src
-
-# stop_queue = asyncio.Queue()
-# url_o = ''
-# filename = ''
-
-
-# def thread_danmaku(loop):
-#     asyncio.set_event_loop(loop)
-#     dmc = src.DanmakuClient(url_o, stop_queue, filename)
-#     loop.run_until_complete(dmc.start())
-#
-#
-# async def danmaku_main(url):
-#     global url_o
-#     url_o = url
-#     executor1 = ThreadPoolExecutor(max_workers=1)
-#     loop = asyncio.new_event_loop()
-#     loop.run_in_executor(executor1, thread_danmaku, loop)
-#     input("输入回车停止运行：\n")
-#     await stop_queue.put(True)
-
-
-class Danmaku(threading.Thread):
-    def __init__(self, filename, url):
-        threading.Thread.__init__(self)
-        self.stop_queue = asyncio.Queue()
-        self.url_ = url
-        self.filename = f"{time.strftime(filename)}.xml"
-
-    def run(self):
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        dmc = src.DanmakuClient(self.url_, self.filename, self.stop_queue)
-        loop.create_task(dmc.start())
-        loop.run_forever()
-
-    async def stop(self):
-        await self.stop_queue.put(True)
-
-
-# 虎牙直播：https://www.huya.com/lpl
-# 斗鱼直播：https://www.douyu.com/9999

biliup/plugins/Danmaku/src/__init__.py

@@ -1,136 +0,0 @@
-import asyncio
-import json
-import os
-import queue
-import re
-import ssl
-import time
-
-import lxml.etree as etree
-import aiofiles
-import aiohttp
-
-from .douyu import Douyu
-from .huya import Huya
-from .bilibili import Bilibili
-
-__all__ = ['DanmakuClient']
-
-
-class DanmakuClient:
-    def __init__(self, url, filename, stop_queue):
-        self.__starttime = time.time()
-        self.__filename = filename
-        self.__url = ''
-        self.__site = None
-        self.__hs = None
-        self.__ws = None
-        self.__stop = False
-        self.__dm_queue = asyncio.Queue()
-        self.__link_status = True
-        self.__stop_queue = stop_queue
-
-        if 'http://' == url[:7] or 'https://' == url[:8]:
-            self.__url = url
-        else:
-            self.__url = 'http://' + url
-        for u, s in {'douyu.com': Douyu, 'huya.com': Huya, 'live.bilibili.com': Bilibili}.items():
-            if re.match(r'^(?:http[s]?://)?.*?%s/(.+?)$' % u, url):
-                self.__site = s
-                self.__u = u
-                break
-        if self.__site is None:
-            print('Invalid link!')
-            exit()
-        self.__hs = aiohttp.ClientSession()
-
-    async def init_ws(self):
-        ws_url, reg_datas = await self.__site.get_ws_info(self.__url)
-        ctx = ssl.create_default_context()
-        ctx.set_ciphers('DEFAULT')
-        self.__ws = await self.__hs.ws_connect(ws_url, ssl_context=ctx)
-        if reg_datas:
-            for reg_data in reg_datas:
-                await self.__ws.send_bytes(reg_data)
-        if not os.path.exists(self.__filename):
-            async with aiofiles.open(self.__filename, mode='w') as f:
-                await f.write("<?xml version='1.0' encoding='UTF-8'?>\n"
-                              "<i xmlns:ns0='http://www.w3.org/1999/XSL/Transform'>\n"
-                              "</i>"
-                              )
-
-    async def heartbeats(self):
-        while not self.__stop and self.__site.heartbeat:
-            await asyncio.sleep(self.__site.heartbeatInterval)
-            try:
-                await self.__ws.send_bytes(self.__site.heartbeat)
-            except Exception as Error:
-                print(f"捕获到异常：{Error}")
-
-    async def fetch_danmaku(self):
-        while not self.__stop:
-            async for msg in self.__ws:
-                self.__link_status = True
-                ms = self.__site.decode_msg(msg.data)
-                for m in ms:
-                    await self.__dm_queue.put(m)
-            await asyncio.sleep(1)
-            await self.init_ws()
-            await asyncio.sleep(1)
-
-    async def print_danmaku(self):
-        parser = etree.XMLParser(recover=True)
-        tree = etree.parse(self.__filename, parser=parser)
-        root = tree.getroot()
-        msg_i = 0
-        msg_col = {'0': '16777215', '1': '16717077', '2': '2000880', '3': '8046667', '4': '16744192', '5': '10172916',
-                   '6': '16738740'}
-        while not self.__stop:
-            try:
-                m = await self.__dm_queue.get()
-            # print(m)
-                if m['msg_type'] == 'danmaku':
-                    d = etree.SubElement(root, 'd')
-                    if 'col' in m:
-                        color = msg_col[m["col"]]
-                    elif 'color' in m:
-                        color = m["color"]
-                    # else:
-                    #     color = '16777215'
-                    msg_time = format(time.time() - self.__starttime, '.3f')
-                    d.set('p', f"{msg_time},1,25,{color},0,0,0,0")
-                    d.text = m["content"]
-
-                    if msg_i >= 5:
-                        etree.indent(root, "\t")
-                        tree.write(self.__filename, encoding="UTF-8", xml_declaration=True, pretty_print=True)
-                        msg_i = 0
-                    else:
-                        msg_i = msg_i + 1
-
-            except Exception as Error:
-                print(f"捕获到异常：{Error}")
-
-    async def start(self):
-        await self.init_ws()
-        await asyncio.gather(
-            self.heartbeats(),
-            self.fetch_danmaku(),
-            self.print_danmaku(),
-            self.stop_check(),
-        )
-
-    async def stop(self):
-        await self.__hs.close()
-        self.__stop = True
-
-    async def stop_check(self):
-        while True:
-            await asyncio.sleep(1)
-            try:
-                stop_tag = self.__stop_queue.get_nowait()
-                if stop_tag:
-                    await self.stop()
-                    break
-            except asyncio.QueueEmpty:
-                pass

biliup/plugins/Danmaku/tars/EndpointF.py

@@ -13,7 +13,7 @@
 # specific language governing permissions and limitations under the License.
 #
 
-from biliup.plugins.Danmaku.src.tars.core import tarscore
+from biliup.plugins.Danmaku.tars.core import tarscore
 
 
 class EndpointF(tarscore.struct):

biliup/plugins/Danmaku/tars/__async.py

@@ -27,9 +27,9 @@ import threading
 
 from urllib3.connectionpool import xrange
 
-from biliup.plugins.Danmaku.src.tars.__logger import tarsLogger
-from biliup.plugins.Danmaku.src.tars.__packet import ResponsePacket
-from biliup.plugins.Danmaku.src.tars.__servantproxy import ServantProxy
+from biliup.plugins.Danmaku.tars.__logger import tarsLogger
+from biliup.plugins.Danmaku.tars.__packet import ResponsePacket
+from biliup.plugins.Danmaku.tars.__servantproxy import ServantProxy
 
 
 class AsyncProcThread:

biliup/plugins/Danmaku/tars/__servantproxy.py

@@ -26,12 +26,12 @@ import threading
 import time
 
 # from __packet import ResponsePacket
-from biliup.plugins.Danmaku.src.tars.__TimeoutQueue import ReqMessage
-from biliup.plugins.Danmaku.src.tars.__logger import tarsLogger
-from biliup.plugins.Danmaku.src.tars.__packet import RequestPacket
-from biliup.plugins.Danmaku.src.tars.__util import util
-from biliup.plugins.Danmaku.src.tars import exception
-from biliup.plugins.Danmaku.src.tars.exception import TarsException
+from biliup.plugins.Danmaku.tars.__TimeoutQueue import ReqMessage
+from biliup.plugins.Danmaku.tars.__logger import tarsLogger
+from biliup.plugins.Danmaku.tars.__packet import RequestPacket
+from biliup.plugins.Danmaku.tars.__util import util
+from biliup.plugins.Danmaku.tars import exception
+from biliup.plugins.Danmaku.tars.exception import TarsException
 
 
 class ServantProxy(object):

biliup/plugins/Danmaku/tars/core.py

@@ -19,10 +19,10 @@
 
 __version__ = "0.0.1"
 
-from biliup.plugins.Danmaku.src.tars.__tars import TarsInputStream
-from biliup.plugins.Danmaku.src.tars.__tars import TarsOutputStream
-from biliup.plugins.Danmaku.src.tars.__tup import TarsUniPacket
-from biliup.plugins.Danmaku.src.tars.__util import util
+from biliup.plugins.Danmaku.tars.__tars import TarsInputStream
+from biliup.plugins.Danmaku.tars.__tars import TarsOutputStream
+from biliup.plugins.Danmaku.tars.__tup import TarsUniPacket
+from biliup.plugins.Danmaku.tars.__util import util
 
 
 class tarscore:

biliup/plugins/bilibili.py

@@ -1,12 +1,12 @@
-import asyncio
+import os
+import random
+import re
 
 import requests
-import os
-import re
-import random
-from . import match1, logger
+
 from biliup.config import config
-from .Danmaku.danmaku_main import Danmaku
+from . import match1, logger
+from biliup.plugins.Danmaku import DanmakuClient
 from ..engine.decorators import Plugin
 from ..engine.download import DownloadBase
 
@@ -132,15 +132,17 @@ class Bilibili(DownloadBase):
                 break
         return True
 
-    def danmaku_download_start(self, filename):
+    async def danmaku_download_start(self, filename):
         if self.bilibili_danmaku:
-            self.danmaku = None
-            self.danmaku = Danmaku(filename, self.url)
-            self.danmaku.start()
+            logger.info("开始弹幕录制")
+            self.danmaku = DanmakuClient(self.url, filename + ".xml")
+            await self.danmaku.start()
 
-    def danmaku_download_stop(self):
+    def close(self):
         if self.bilibili_danmaku:
-            asyncio.run(self.danmaku.stop())
+            self.danmaku.stop()
+            logger.info("结束弹幕录制")
+            # asyncio.run(self.danmaku.stop())
 
 
 def get_play_info(s, custom_api, official_api_host, params):

biliup/plugins/douyu.py

@@ -1,18 +1,13 @@
-import asyncio
-import subprocess
-import sys
 from collections import namedtuple
-import time
-from urllib.parse import urlparse
 
 import requests
-
 from ykdl.util.match import match1
+
 from biliup.config import config
+from biliup.plugins.Danmaku import DanmakuClient
 from ..engine.decorators import Plugin
-from ..engine.download import DownloadBase, get_valid_filename, stream_gears_download
+from ..engine.download import DownloadBase
 from ..plugins import logger
-from .Danmaku.danmaku_main import Danmaku
 
 
 @Plugin.download(regexp=r'(?:https?://)?(?:(?:www|m)\.)?douyu\.com')
@@ -62,12 +57,13 @@ class Douyu(DownloadBase):
             self.room_title = roominfo['room_name']
             return True
 
-    def danmaku_download_start(self, filename):
+    async def danmaku_download_start(self, filename):
         if self.douyu_danmaku:
-            self.danmaku = None
-            self.danmaku = Danmaku(filename, self.url)
-            self.danmaku.start()
+            logger.info("开始弹幕录制")
+            self.danmaku = DanmakuClient(self.url, filename + ".xml")
+            await self.danmaku.start()
 
-    def danmaku_download_stop(self):
+    def close(self):
         if self.douyu_danmaku:
-            asyncio.run(self.danmaku.stop())
+            self.danmaku.stop()
+            logger.info("结束弹幕录制")

biliup/plugins/huya.py

@@ -1,4 +1,3 @@
-import asyncio
 import base64
 import html
 import json
@@ -6,10 +5,10 @@ import json
 import requests
 
 from biliup.config import config
-from .Danmaku.danmaku_main import Danmaku
+from biliup.plugins.Danmaku import DanmakuClient
 from ..engine.decorators import Plugin
-from ..plugins import match1, logger
 from ..engine.download import DownloadBase
+from ..plugins import match1, logger
 
 
 @Plugin.download(regexp=r'(?:https?://)?(?:(?:www|m)\.)?huya\.com')
@@ -53,12 +52,13 @@ class Huya(DownloadBase):
             self.room_title = json.loads(huya)['data'][0]['gameLiveInfo']['introduction']
             return True
 
-    def danmaku_download_start(self, filename):
+    async def danmaku_download_start(self, filename):
         if self.huya_danmaku:
-            self.danmaku = None
-            self.danmaku = Danmaku(filename, self.url)
-            self.danmaku.start()
+            logger.info("开始弹幕录制")
+            self.danmaku = DanmakuClient(self.url, filename + ".xml")
+            await self.danmaku.start()
 
-    def danmaku_download_stop(self):
+    def close(self):
         if self.huya_danmaku:
-            asyncio.run(self.danmaku.stop())
+            self.danmaku.stop()
+            logger.info("结束弹幕录制")