From fc9a9d2386af25b3f9719df405a2efba95472033 Mon Sep 17 00:00:00 2001
From: RockChinQ <1010553892@qq.com>
Date: Tue, 2 Apr 2024 22:33:06 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E7=BC=BA=E5=A4=B1=E7=9A=84=20psutil=20?=
 =?UTF-8?q?=E4=BE=9D=E8=B5=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.py                    | 25 ++++++++++++++-----------
 pkg/core/bootutils/deps.py |  1 +
 requirements.txt           |  3 ++-
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/main.py b/main.py
index 055a530a..d7637b18 100644
--- a/main.py
+++ b/main.py
@@ -32,6 +32,20 @@ async def main_entry():
         print("已自动安装缺失的依赖包，请重启程序。")
         sys.exit(0)
 
+    # 检查命令行
+    import os
+
+    if os.name == 'nt':
+        import psutil
+        allowed_parent_process = ['cmd.exe', 'powershell.exe', 'wsl.exe']
+
+        parent_process = psutil.Process(os.getppid()).name()
+
+        if parent_process not in allowed_parent_process:
+            print("请在命令行中运行此程序。")
+            input("按任意键退出...")
+            exit(0)
+
     # 检查配置文件
 
     from pkg.core.bootutils import files
@@ -51,17 +65,6 @@ async def main_entry():
 
 if __name__ == '__main__':
     import os
-    import psutil
-
-    if os.name == 'nt':
-        allowed_parent_process = ['cmd.exe', 'powershell.exe', 'wsl.exe']
-
-        parent_process = psutil.Process(os.getppid()).name()
-
-        if parent_process not in allowed_parent_process:
-            print("请在命令行中运行此程序。")
-            input("按任意键退出...")
-            exit(0)
 
     # 检查本目录是否有main.py，且包含QChatGPT字符串
     invalid_pwd = False
diff --git a/pkg/core/bootutils/deps.py b/pkg/core/bootutils/deps.py
index 3b44e9cc..4adf1323 100644
--- a/pkg/core/bootutils/deps.py
+++ b/pkg/core/bootutils/deps.py
@@ -13,6 +13,7 @@
     "tiktoken": "tiktoken",
     "yaml": "pyyaml",
     "aiohttp": "aiohttp",
+    "psutil": "psutil",
 }
 
 
diff --git a/requirements.txt b/requirements.txt
index 28c0ecb6..f04bdc9a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,4 +12,5 @@ PyYaml
 aiohttp
 pydantic
 websockets
-urllib3
\ No newline at end of file
+urllib3
+psutil
\ No newline at end of file