From 4f66fae3fb38faf203332c7075051e93b551cc20 Mon Sep 17 00:00:00 2001
From: Kyomotoi <0w0@imki.moe>
Date: Thu, 30 Mar 2023 15:13:39 +0800
Subject: =?UTF-8?q?=F0=9F=8E=A8=20=E4=BC=98=E5=8C=96=E5=B7=A5=E5=85=B7?=
 =?UTF-8?q?=E7=B1=BB=E6=80=A7=E8=83=BD=E8=AF=BB=E5=8F=96=E9=83=A8=E5=88=86?=
 =?UTF-8?q?=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ATRI/utils/machine.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

(limited to 'ATRI')

diff --git a/ATRI/utils/machine.py b/ATRI/utils/machine.py
index a2b5b55..ee2b533 100644
--- a/ATRI/utils/machine.py
+++ b/ATRI/utils/machine.py
@@ -1,13 +1,12 @@
 import psutil
 import platform
 from pydantic import BaseModel
-from sys import platform as pf, modules
-from importlib.util import find_spec
+from sys import platform as pf
 
 from .apscheduler import scheduler
 
-if pf == "win32" and find_spec("win32com"):
-    import wmi  # type: ignore
+if pf == "win32":
+    import wmi
     from win32com.client import GetObject
 
 
@@ -23,6 +22,7 @@ class CpuInfo(BaseModel):
     max_freq: str
     current_freq: str
     percent: float
+    process: int
 
 
 class MemInfo(BaseModel):
@@ -45,6 +45,7 @@ class DiskInfo(BaseModel):
     total: int
     used: int
     free: int
+    drive: int
 
 
 class NetInfo(BaseModel):
@@ -63,16 +64,17 @@ def get_platform_info() -> PlatformInfo:
 
 def get_cpu_info() -> CpuInfo:
     cpu_name = platform.processor()
-    if pf == "win32" and "win32com" in modules:
+    if pf == "win32":
         winm = GetObject("winmgmts:root\cimv2")
         cpus = winm.ExecQuery("SELECT * FROM Win32_Processor")
-        cpu_name = cpus[0].Name
+        cpu_name = cpus[0].Name.strip()
 
     cpu_count = psutil.cpu_count(False)
     _freq = psutil.cpu_freq()
     cpu_max_freq = f"{'%.2f'%(_freq.max / 1000)}"
-    cpu_current_freq = f"{'%.2f'%(_freq.current)}"
+    cpu_current_freq = f"{'%.2f'%(_freq.current / 1000)}"
     cpu_percent = psutil.cpu_percent(interval=0.1)
+    process = len(psutil.pids())
 
     return CpuInfo(
         name=cpu_name,
@@ -80,6 +82,7 @@ def get_cpu_info() -> CpuInfo:
         max_freq=cpu_max_freq,
         current_freq=cpu_current_freq,
         percent=cpu_percent,
+        process=process,
     )
 
 
@@ -106,7 +109,8 @@ def get_disk_info():
     disk_total = int()
     disk_used = int()
     disk_free = int()
-    if pf == "win32" and "win32com" in modules:
+    disk_drive = 1
+    if pf == "win32":
         w = wmi.WMI()
         disk_list = [d.DeviceID for d in w.Win32_LogicalDisk()]
         for i in disk_list:
@@ -116,6 +120,7 @@ def get_disk_info():
             disk_free += disk.free
 
         disk_name = w.Win32_DiskDrive()[0].Model
+        disk_drive = len(w.Win32_DiskDrive())
     else:
         disk = psutil.disk_usage("/")
         disk_total = disk.total
@@ -127,6 +132,7 @@ def get_disk_info():
         total=disk_total,
         used=disk_used,
         free=disk_free,
+        drive=disk_drive,
     )
 
 
-- 
cgit v1.2.3