Преглед изворни кода

feat: Add configuration items to support custom data directories and facilitate the storage of itchat.pkl

李超 пре 3 година
родитељ
комит
29fbf69945
2 измењених фајлова са 15 додато и 6 уклоњено
  1. 4 3
      channel/wechat/wechat_channel.py
  2. 11 3
      config.py

+ 4 - 3
channel/wechat/wechat_channel.py

@@ -18,7 +18,7 @@ from lib import itchat
 from lib.itchat.content import *
 from bridge.reply import *
 from bridge.context import *
-from config import conf
+from config import conf, get_data_path
 from common.time_check import time_checker
 from common.expired_dict import ExpiredDict
 from plugins import *
@@ -97,13 +97,14 @@ class WechatChannel(ChatChannel):
         itchat.instance.receivingRetryCount = 600  # 修改断线超时时间
         # login by scan QRCode
         hotReload = conf().get('hot_reload', False)
+        status_path = os.path.join(get_data_path(), "itchat.pkl")
         try:
-            itchat.auto_login(enableCmdQR=2, hotReload=hotReload, qrCallback=qrCallback)
+            itchat.auto_login(enableCmdQR=2, hotReload=hotReload, statusStorageDir=status_path, qrCallback=qrCallback)
         except Exception as e:
             if hotReload:
                 logger.error("Hot reload failed, try to login without hot reload")
                 itchat.logout()
-                os.remove("itchat.pkl")
+                os.remove(status_path)
                 itchat.auto_login(enableCmdQR=2, hotReload=hotReload, qrCallback=qrCallback)
             else:
                 raise e

+ 11 - 3
config.py

@@ -93,6 +93,8 @@ available_setting = {
 
     "debug": False,  # 是否开启debug模式,开启后会打印更多日志
 
+    "config_data_path": "",    # 数据目录
+
     # 插件配置
     "plugin_trigger_prefix": "$",  # 规范插件提供聊天相关指令的前缀,建议不要和管理员指令前缀"#"冲突
 }
@@ -130,7 +132,7 @@ class Config(dict):
 
     def load_user_datas(self):
         try:
-            with open('user_datas.pkl', 'rb') as f:
+            with open(os.path.join(get_data_path(), 'user_datas.pkl'), 'rb') as f:
                 self.user_datas = pickle.load(f)
                 logger.info("[Config] User datas loaded.")
         except FileNotFoundError as e:
@@ -141,7 +143,7 @@ class Config(dict):
 
     def save_user_datas(self):
         try:
-            with open('user_datas.pkl', 'wb') as f:
+            with open(os.path.join(get_data_path(), 'user_datas.pkl'), 'wb') as f:
                 pickle.dump(self.user_datas, f)
                 logger.info("[Config] User datas saved.")
         except Exception as e:
@@ -196,6 +198,12 @@ def read_file(path):
     with open(path, mode='r', encoding='utf-8') as f:
         return f.read()
 
-
 def conf():
     return config
+
+def get_data_path():
+    data_path = os.path.join(get_root(), conf().get('config_data_path', ""))
+    if not os.path.exists(data_path):
+        logger.info("[INIT] data path not exists, create it: {}".format(data_path))
+        os.makedirs(data_path)
+    return data_path