formatting code

channel/chat_channel.py

@@ -1,137 +1,172 @@
-
-
-from asyncio import CancelledError
-from concurrent.futures import Future, ThreadPoolExecutor
 import os
 import re
 import threading
 import time
-from common.dequeue import Dequeue
-from channel.channel import Channel
-from bridge.reply import *
+from asyncio import CancelledError
+from concurrent.futures import Future, ThreadPoolExecutor
+
 from bridge.context import *
-from config import conf
+from bridge.reply import *
+from channel.channel import Channel
+from common.dequeue import Dequeue
 from common.log import logger
+from config import conf
 from plugins import *
+
 try:
     from voice.audio_convert import any_to_wav
 except Exception as e:
     pass
 
+
 # 抽象类, 它包含了与消息通道无关的通用处理逻辑
 class ChatChannel(Channel):
-    name = None # 登录的用户名
-    user_id = None # 登录的用户id
-    futures = {} # 记录每个session_id提交到线程池的future对象, 用于重置会话时把没执行的future取消掉，正在执行的不会被取消
-    sessions = {} # 用于控制并发，每个session_id同时只能有一个context在处理
-    lock = threading.Lock() # 用于控制对sessions的访问
+    name = None  # 登录的用户名
+    user_id = None  # 登录的用户id
+    futures = {}  # 记录每个session_id提交到线程池的future对象, 用于重置会话时把没执行的future取消掉，正在执行的不会被取消
+    sessions = {}  # 用于控制并发，每个session_id同时只能有一个context在处理
+    lock = threading.Lock()  # 用于控制对sessions的访问
     handler_pool = ThreadPoolExecutor(max_workers=8)  # 处理消息的线程池
 
     def __init__(self):
         _thread = threading.Thread(target=self.consume)
         _thread.setDaemon(True)
         _thread.start()
-        
 
     # 根据消息构造context，消息内容相关的触发项写在这里
     def _compose_context(self, ctype: ContextType, content, **kwargs):
         context = Context(ctype, content)
         context.kwargs = kwargs
-        # context首次传入时，origin_ctype是None, 
+        # context首次传入时，origin_ctype是None,
         # 引入的起因是：当输入语音时，会嵌套生成两个context，第一步语音转文本，第二步通过文本生成文字回复。
         # origin_ctype用于第二步文本回复时，判断是否需要匹配前缀，如果是私聊的语音，就不需要匹配前缀
-        if 'origin_ctype' not in context:  
-            context['origin_ctype'] = ctype
+        if "origin_ctype" not in context:
+            context["origin_ctype"] = ctype
         # context首次传入时，receiver是None，根据类型设置receiver
-        first_in = 'receiver' not in context
+        first_in = "receiver" not in context
         # 群名匹配过程，设置session_id和receiver
-        if first_in: # context首次传入时，receiver是None，根据类型设置receiver
+        if first_in:  # context首次传入时，receiver是None，根据类型设置receiver
             config = conf()
-            cmsg = context['msg']
+            cmsg = context["msg"]
             if context.get("isgroup", False):
                 group_name = cmsg.other_user_nickname
                 group_id = cmsg.other_user_id
 
-                group_name_white_list = config.get('group_name_white_list', [])
-                group_name_keyword_white_list = config.get('group_name_keyword_white_list', [])
-                if any([group_name in group_name_white_list, 'ALL_GROUP' in group_name_white_list, check_contain(group_name, group_name_keyword_white_list)]):
-                    group_chat_in_one_session = conf().get('group_chat_in_one_session', [])
+                group_name_white_list = config.get("group_name_white_list", [])
+                group_name_keyword_white_list = config.get(
+                    "group_name_keyword_white_list", []
+                )
+                if any(
+                    [
+                        group_name in group_name_white_list,
+                        "ALL_GROUP" in group_name_white_list,
+                        check_contain(group_name, group_name_keyword_white_list),
+                    ]
+                ):
+                    group_chat_in_one_session = conf().get(
+                        "group_chat_in_one_session", []
+                    )
                     session_id = cmsg.actual_user_id
-                    if any([group_name in group_chat_in_one_session, 'ALL_GROUP' in group_chat_in_one_session]):
+                    if any(
+                        [
+                            group_name in group_chat_in_one_session,
+                            "ALL_GROUP" in group_chat_in_one_session,
+                        ]
+                    ):
                         session_id = group_id
                 else:
                     return None
-                context['session_id'] = session_id
-                context['receiver'] = group_id
+                context["session_id"] = session_id
+                context["receiver"] = group_id
             else:
-                context['session_id'] = cmsg.other_user_id
-                context['receiver'] = cmsg.other_user_id
-            e_context = PluginManager().emit_event(EventContext(Event.ON_RECEIVE_MESSAGE, {'channel': self, 'context': context}))
-            context = e_context['context']
+                context["session_id"] = cmsg.other_user_id
+                context["receiver"] = cmsg.other_user_id
+            e_context = PluginManager().emit_event(
+                EventContext(
+                    Event.ON_RECEIVE_MESSAGE, {"channel": self, "context": context}
+                )
+            )
+            context = e_context["context"]
             if e_context.is_pass() or context is None:
                 return context
-            if cmsg.from_user_id == self.user_id and not config.get('trigger_by_self', True):
+            if cmsg.from_user_id == self.user_id and not config.get(
+                "trigger_by_self", True
+            ):
                 logger.debug("[WX]self message skipped")
                 return None
 
         # 消息内容匹配过程，并处理content
         if ctype == ContextType.TEXT:
-            if first_in and "」\n- - - - - - -" in content: # 初次匹配 过滤引用消息
+            if first_in and "」\n- - - - - - -" in content:  # 初次匹配 过滤引用消息
                 logger.debug("[WX]reference query skipped")
                 return None
-            
-            if context.get("isgroup", False): # 群聊
+
+            if context.get("isgroup", False):  # 群聊
                 # 校验关键字
-                match_prefix = check_prefix(content, conf().get('group_chat_prefix'))
-                match_contain = check_contain(content, conf().get('group_chat_keyword'))
+                match_prefix = check_prefix(content, conf().get("group_chat_prefix"))
+                match_contain = check_contain(content, conf().get("group_chat_keyword"))
                 flag = False
                 if match_prefix is not None or match_contain is not None:
                     flag = True
                     if match_prefix:
-                        content = content.replace(match_prefix, '', 1).strip()
-                if context['msg'].is_at:
+                        content = content.replace(match_prefix, "", 1).strip()
+                if context["msg"].is_at:
                     logger.info("[WX]receive group at")
                     if not conf().get("group_at_off", False):
                         flag = True
-                    pattern = f'@{self.name}(\u2005|\u0020)'
-                    content = re.sub(pattern, r'', content)
-                
+                    pattern = f"@{self.name}(\u2005|\u0020)"
+                    content = re.sub(pattern, r"", content)
+
                 if not flag:
                     if context["origin_ctype"] == ContextType.VOICE:
-                        logger.info("[WX]receive group voice, but checkprefix didn't match")
+                        logger.info(
+                            "[WX]receive group voice, but checkprefix didn't match"
+                        )
                     return None
-            else: # 单聊
-                match_prefix = check_prefix(content, conf().get('single_chat_prefix',['']))
-                if match_prefix is not None: # 判断如果匹配到自定义前缀，则返回过滤掉前缀+空格后的内容
-                    content = content.replace(match_prefix, '', 1).strip()
-                elif context["origin_ctype"] == ContextType.VOICE: # 如果源消息是私聊的语音消息，允许不匹配前缀，放宽条件
+            else:  # 单聊
+                match_prefix = check_prefix(
+                    content, conf().get("single_chat_prefix", [""])
+                )
+                if match_prefix is not None:  # 判断如果匹配到自定义前缀，则返回过滤掉前缀+空格后的内容
+                    content = content.replace(match_prefix, "", 1).strip()
+                elif (
+                    context["origin_ctype"] == ContextType.VOICE
+                ):  # 如果源消息是私聊的语音消息，允许不匹配前缀，放宽条件
                     pass
                 else:
-                    return None     
-                                                  
-            img_match_prefix = check_prefix(content, conf().get('image_create_prefix'))
+                    return None
+
+            img_match_prefix = check_prefix(content, conf().get("image_create_prefix"))
             if img_match_prefix:
-                content = content.replace(img_match_prefix, '', 1)
+                content = content.replace(img_match_prefix, "", 1)
                 context.type = ContextType.IMAGE_CREATE
             else:
                 context.type = ContextType.TEXT
             context.content = content.strip()
-            if 'desire_rtype' not in context and conf().get('always_reply_voice') and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE:
-                context['desire_rtype'] = ReplyType.VOICE
-        elif context.type == ContextType.VOICE: 
-            if 'desire_rtype' not in context and conf().get('voice_reply_voice') and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE:
-                context['desire_rtype'] = ReplyType.VOICE
+            if (
+                "desire_rtype" not in context
+                and conf().get("always_reply_voice")
+                and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE
+            ):
+                context["desire_rtype"] = ReplyType.VOICE
+        elif context.type == ContextType.VOICE:
+            if (
+                "desire_rtype" not in context
+                and conf().get("voice_reply_voice")
+                and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE
+            ):
+                context["desire_rtype"] = ReplyType.VOICE
 
         return context
 
     def _handle(self, context: Context):
         if context is None or not context.content:
             return
-        logger.debug('[WX] ready to handle context: {}'.format(context))
+        logger.debug("[WX] ready to handle context: {}".format(context))
         # reply的构建步骤
         reply = self._generate_reply(context)
 
-        logger.debug('[WX] ready to decorate reply: {}'.format(reply))
+        logger.debug("[WX] ready to decorate reply: {}".format(reply))
         # reply的包装步骤
         reply = self._decorate_reply(context, reply)
 
@@ -139,20 +174,31 @@ class ChatChannel(Channel):
         self._send_reply(context, reply)
 
     def _generate_reply(self, context: Context, reply: Reply = Reply()) -> Reply:
-        e_context = PluginManager().emit_event(EventContext(Event.ON_HANDLE_CONTEXT, {
-            'channel': self, 'context': context, 'reply': reply}))
-        reply = e_context['reply']
+        e_context = PluginManager().emit_event(
+            EventContext(
+                Event.ON_HANDLE_CONTEXT,
+                {"channel": self, "context": context, "reply": reply},
+            )
+        )
+        reply = e_context["reply"]
         if not e_context.is_pass():
-            logger.debug('[WX] ready to handle context: type={}, content={}'.format(context.type, context.content))
-            if context.type == ContextType.TEXT or context.type == ContextType.IMAGE_CREATE:  # 文字和图片消息
+            logger.debug(
+                "[WX] ready to handle context: type={}, content={}".format(
+                    context.type, context.content
+                )
+            )
+            if (
+                context.type == ContextType.TEXT
+                or context.type == ContextType.IMAGE_CREATE
+            ):  # 文字和图片消息
                 reply = super().build_reply_content(context.content, context)
             elif context.type == ContextType.VOICE:  # 语音消息
-                cmsg = context['msg']
+                cmsg = context["msg"]
                 cmsg.prepare()
                 file_path = context.content
-                wav_path = os.path.splitext(file_path)[0] + '.wav'
+                wav_path = os.path.splitext(file_path)[0] + ".wav"
                 try:
-                    any_to_wav(file_path, wav_path) 
+                    any_to_wav(file_path, wav_path)
                 except Exception as e:  # 转换失败，直接使用mp3，对于某些api，mp3也可以识别
                     logger.warning("[WX]any to wav error, use raw path. " + str(e))
                     wav_path = file_path
@@ -169,7 +215,8 @@ class ChatChannel(Channel):
 
                 if reply.type == ReplyType.TEXT:
                     new_context = self._compose_context(
-                        ContextType.TEXT, reply.content, **context.kwargs)
+                        ContextType.TEXT, reply.content, **context.kwargs
+                    )
                     if new_context:
                         reply = self._generate_reply(new_context)
                     else:
@@ -177,18 +224,21 @@ class ChatChannel(Channel):
             elif context.type == ContextType.IMAGE:  # 图片消息，当前无默认逻辑
                 pass
             else:
-                logger.error('[WX] unknown context type: {}'.format(context.type))
+                logger.error("[WX] unknown context type: {}".format(context.type))
                 return
         return reply
 
     def _decorate_reply(self, context: Context, reply: Reply) -> Reply:
         if reply and reply.type:
-            e_context = PluginManager().emit_event(EventContext(Event.ON_DECORATE_REPLY, {
-                'channel': self, 'context': context, 'reply': reply}))
-            reply = e_context['reply']
-            desire_rtype = context.get('desire_rtype')
+            e_context = PluginManager().emit_event(
+                EventContext(
+                    Event.ON_DECORATE_REPLY,
+                    {"channel": self, "context": context, "reply": reply},
+                )
+            )
+            reply = e_context["reply"]
+            desire_rtype = context.get("desire_rtype")
             if not e_context.is_pass() and reply and reply.type:
-                
                 if reply.type in self.NOT_SUPPORT_REPLYTYPE:
                     logger.error("[WX]reply type not support: " + str(reply.type))
                     reply.type = ReplyType.ERROR
@@ -196,59 +246,91 @@ class ChatChannel(Channel):
 
                 if reply.type == ReplyType.TEXT:
                     reply_text = reply.content
-                    if desire_rtype == ReplyType.VOICE and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE:
+                    if (
+                        desire_rtype == ReplyType.VOICE
+                        and ReplyType.VOICE not in self.NOT_SUPPORT_REPLYTYPE
+                    ):
                         reply = super().build_text_to_voice(reply.content)
                         return self._decorate_reply(context, reply)
                     if context.get("isgroup", False):
-                        reply_text = '@' +  context['msg'].actual_user_nickname + ' ' + reply_text.strip()
-                        reply_text = conf().get("group_chat_reply_prefix", "") + reply_text
+                        reply_text = (
+                            "@"
+                            + context["msg"].actual_user_nickname
+                            + " "
+                            + reply_text.strip()
+                        )
+                        reply_text = (
+                            conf().get("group_chat_reply_prefix", "") + reply_text
+                        )
                     else:
-                        reply_text = conf().get("single_chat_reply_prefix", "") + reply_text
+                        reply_text = (
+                            conf().get("single_chat_reply_prefix", "") + reply_text
+                        )
                     reply.content = reply_text
                 elif reply.type == ReplyType.ERROR or reply.type == ReplyType.INFO:
-                    reply.content = "["+str(reply.type)+"]\n" + reply.content
-                elif reply.type == ReplyType.IMAGE_URL or reply.type == ReplyType.VOICE or reply.type == ReplyType.IMAGE:
+                    reply.content = "[" + str(reply.type) + "]\n" + reply.content
+                elif (
+                    reply.type == ReplyType.IMAGE_URL
+                    or reply.type == ReplyType.VOICE
+                    or reply.type == ReplyType.IMAGE
+                ):
                     pass
                 else:
-                    logger.error('[WX] unknown reply type: {}'.format(reply.type))
+                    logger.error("[WX] unknown reply type: {}".format(reply.type))
                     return
-            if desire_rtype and desire_rtype != reply.type and reply.type not in [ReplyType.ERROR, ReplyType.INFO]:
-                logger.warning('[WX] desire_rtype: {}, but reply type: {}'.format(context.get('desire_rtype'), reply.type))
+            if (
+                desire_rtype
+                and desire_rtype != reply.type
+                and reply.type not in [ReplyType.ERROR, ReplyType.INFO]
+            ):
+                logger.warning(
+                    "[WX] desire_rtype: {}, but reply type: {}".format(
+                        context.get("desire_rtype"), reply.type
+                    )
+                )
             return reply
 
     def _send_reply(self, context: Context, reply: Reply):
         if reply and reply.type:
-            e_context = PluginManager().emit_event(EventContext(Event.ON_SEND_REPLY, {
-                'channel': self, 'context': context, 'reply': reply}))
-            reply = e_context['reply']
+            e_context = PluginManager().emit_event(
+                EventContext(
+                    Event.ON_SEND_REPLY,
+                    {"channel": self, "context": context, "reply": reply},
+                )
+            )
+            reply = e_context["reply"]
             if not e_context.is_pass() and reply and reply.type:
-                logger.debug('[WX] ready to send reply: {}, context: {}'.format(reply, context))
+                logger.debug(
+                    "[WX] ready to send reply: {}, context: {}".format(reply, context)
+                )
                 self._send(reply, context)
 
-    def _send(self, reply: Reply, context: Context, retry_cnt = 0):
+    def _send(self, reply: Reply, context: Context, retry_cnt=0):
         try:
             self.send(reply, context)
         except Exception as e:
-            logger.error('[WX] sendMsg error: {}'.format(str(e)))
+            logger.error("[WX] sendMsg error: {}".format(str(e)))
             if isinstance(e, NotImplementedError):
                 return
             logger.exception(e)
             if retry_cnt < 2:
-                time.sleep(3+3*retry_cnt)
-                self._send(reply, context, retry_cnt+1)
+                time.sleep(3 + 3 * retry_cnt)
+                self._send(reply, context, retry_cnt + 1)
 
-    def _success_callback(self, session_id, **kwargs):# 线程正常结束时的回调函数
+    def _success_callback(self, session_id, **kwargs):  # 线程正常结束时的回调函数
         logger.debug("Worker return success, session_id = {}".format(session_id))
 
-    def _fail_callback(self, session_id, exception, **kwargs): # 线程异常结束时的回调函数
+    def _fail_callback(self, session_id, exception, **kwargs):  # 线程异常结束时的回调函数
         logger.exception("Worker return exception: {}".format(exception))
 
     def _thread_pool_callback(self, session_id, **kwargs):
-        def func(worker:Future):
+        def func(worker: Future):
             try:
                 worker_exception = worker.exception()
                 if worker_exception:
-                    self._fail_callback(session_id, exception = worker_exception, **kwargs)
+                    self._fail_callback(
+                        session_id, exception=worker_exception, **kwargs
+                    )
                 else:
                     self._success_callback(session_id, **kwargs)
             except CancelledError as e:
@@ -257,15 +339,19 @@ class ChatChannel(Channel):
                 logger.exception("Worker raise exception: {}".format(e))
             with self.lock:
                 self.sessions[session_id][1].release()
+
         return func
 
     def produce(self, context: Context):
-        session_id = context['session_id']
+        session_id = context["session_id"]
         with self.lock:
             if session_id not in self.sessions:
-                self.sessions[session_id] = [Dequeue(), threading.BoundedSemaphore(conf().get("concurrency_in_session", 4))]
-            if context.type == ContextType.TEXT and context.content.startswith("#"): 
-                self.sessions[session_id][0].putleft(context) # 优先处理管理命令
+                self.sessions[session_id] = [
+                    Dequeue(),
+                    threading.BoundedSemaphore(conf().get("concurrency_in_session", 4)),
+                ]
+            if context.type == ContextType.TEXT and context.content.startswith("#"):
+                self.sessions[session_id][0].putleft(context)  # 优先处理管理命令
             else:
                 self.sessions[session_id][0].put(context)
 
@@ -276,44 +362,58 @@ class ChatChannel(Channel):
                 session_ids = list(self.sessions.keys())
                 for session_id in session_ids:
                     context_queue, semaphore = self.sessions[session_id]
-                    if semaphore.acquire(blocking = False): # 等线程处理完毕才能删除
+                    if semaphore.acquire(blocking=False):  # 等线程处理完毕才能删除
                         if not context_queue.empty():
                             context = context_queue.get()
                             logger.debug("[WX] consume context: {}".format(context))
-                            future:Future = self.handler_pool.submit(self._handle, context)
-                            future.add_done_callback(self._thread_pool_callback(session_id, context = context))
+                            future: Future = self.handler_pool.submit(
+                                self._handle, context
+                            )
+                            future.add_done_callback(
+                                self._thread_pool_callback(session_id, context=context)
+                            )
                             if session_id not in self.futures:
                                 self.futures[session_id] = []
                             self.futures[session_id].append(future)
-                        elif semaphore._initial_value == semaphore._value+1: # 除了当前，没有任务再申请到信号量，说明所有任务都处理完毕
-                            self.futures[session_id] = [t for t in self.futures[session_id] if not t.done()]
-                            assert len(self.futures[session_id]) == 0, "thread pool error"
+                        elif (
+                            semaphore._initial_value == semaphore._value + 1
+                        ):  # 除了当前，没有任务再申请到信号量，说明所有任务都处理完毕
+                            self.futures[session_id] = [
+                                t for t in self.futures[session_id] if not t.done()
+                            ]
+                            assert (
+                                len(self.futures[session_id]) == 0
+                            ), "thread pool error"
                             del self.sessions[session_id]
                         else:
                             semaphore.release()
             time.sleep(0.1)
 
     # 取消session_id对应的所有任务，只能取消排队的消息和已提交线程池但未执行的任务
-    def cancel_session(self, session_id): 
+    def cancel_session(self, session_id):
         with self.lock:
             if session_id in self.sessions:
                 for future in self.futures[session_id]:
                     future.cancel()
                 cnt = self.sessions[session_id][0].qsize()
-                if cnt>0:
-                    logger.info("Cancel {} messages in session {}".format(cnt, session_id))
+                if cnt > 0:
+                    logger.info(
+                        "Cancel {} messages in session {}".format(cnt, session_id)
+                    )
                 self.sessions[session_id][0] = Dequeue()
-    
+
     def cancel_all_session(self):
         with self.lock:
             for session_id in self.sessions:
                 for future in self.futures[session_id]:
                     future.cancel()
                 cnt = self.sessions[session_id][0].qsize()
-                if cnt>0:
-                    logger.info("Cancel {} messages in session {}".format(cnt, session_id))
+                if cnt > 0:
+                    logger.info(
+                        "Cancel {} messages in session {}".format(cnt, session_id)
+                    )
                 self.sessions[session_id][0] = Dequeue()
-    
+
 
 def check_prefix(content, prefix_list):
     if not prefix_list:
@@ -323,6 +423,7 @@ def check_prefix(content, prefix_list):
             return prefix
     return None
 
+
 def check_contain(content, keyword_list):
     if not keyword_list:
         return None