From 7708eaa82c7dee855f9d577a64274aecd86730a1 Mon Sep 17 00:00:00 2001
From: RockChinQ <1010553892@qq.com>
Date: Sun, 26 Nov 2023 17:33:13 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=B8=BA=20context.py=20=E4=B8=AD?=
 =?UTF-8?q?=E7=9A=84=E6=96=B9=E6=B3=95=E6=B7=BB=E5=8A=A0=E7=B1=BB=E5=9E=8B?=
 =?UTF-8?q?=E6=8F=90=E7=A4=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/utils/context.py | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/pkg/utils/context.py b/pkg/utils/context.py
index 0da18228..b208dac8 100644
--- a/pkg/utils/context.py
+++ b/pkg/utils/context.py
@@ -1,6 +1,13 @@
+from __future__ import annotations
+
 import threading
 from . import threadctl
 
+from ..database import manager as db_mgr
+from ..openai import manager as openai_mgr
+from ..qqbot import manager as qqbot_mgr
+from ..plugin import host as plugin_host
+
 
 context = {
     'inst': {
@@ -29,59 +36,59 @@ def get_config():
     return t
 
 
-def set_database_manager(inst):
+def set_database_manager(inst: db_mgr.DatabaseManager):
     context_lock.acquire()
     context['inst']['database.manager.DatabaseManager'] = inst
     context_lock.release()
 
 
-def get_database_manager():
+def get_database_manager() -> db_mgr.DatabaseManager:
     context_lock.acquire()
     t = context['inst']['database.manager.DatabaseManager']
     context_lock.release()
     return t
 
 
-def set_openai_manager(inst):
+def set_openai_manager(inst: openai_mgr.OpenAIInteract):
     context_lock.acquire()
     context['inst']['openai.manager.OpenAIInteract'] = inst
     context_lock.release()
 
 
-def get_openai_manager():
+def get_openai_manager() -> openai_mgr.OpenAIInteract:
     context_lock.acquire()
     t = context['inst']['openai.manager.OpenAIInteract']
     context_lock.release()
     return t
 
 
-def set_qqbot_manager(inst):
+def set_qqbot_manager(inst: qqbot_mgr.QQBotManager):
     context_lock.acquire()
     context['inst']['qqbot.manager.QQBotManager'] = inst
     context_lock.release()
 
 
-def get_qqbot_manager():
+def get_qqbot_manager() -> qqbot_mgr.QQBotManager:
     context_lock.acquire()
     t = context['inst']['qqbot.manager.QQBotManager']
     context_lock.release()
     return t
 
 
-def set_plugin_host(inst):
+def set_plugin_host(inst: plugin_host.PluginHost):
     context_lock.acquire()
     context['plugin_host'] = inst
     context_lock.release()
 
 
-def get_plugin_host():
+def get_plugin_host() -> plugin_host.PluginHost:
     context_lock.acquire()
     t = context['plugin_host']
     context_lock.release()
     return t
 
 
-def set_thread_ctl(inst):
+def set_thread_ctl(inst: threadctl.ThreadCtl):
     context_lock.acquire()
     context['pool_ctl'] = inst
     context_lock.release()