0%

自己使用curses写一个类似于htop的命令行界面程序

发表于 2024-07-31 阅读次数：
本文字数： 7k 阅读时长 ≈ 13 分钟

在命令行循环显示一个内容而不向下叠加

使用curses.wrapper包裹工作函数
工作函数需要接受一个参数stdscr
stdscr的用法
stdscr.addstr(y, x, str, curses.color_pair)

第一个参数是行数，第二个是列数，第三个是颜色对

# 定义
curses.start_color()
curses.init_pair(1, curses.COLOR_WHITE, curses.COLOR_BLACK)
bold = curses.A_BOLD
# 使用
stdscr.addstr(y, x, "****", curses.color_pair(1) | bold)

获取CPU核GPU信息

使用psutil和pynvml库
具体参考如下用例

用psutil的process_iter获取CPU占用排名靠前的进程

def get_top_cpu_processes():
    processes = psutil.process_iter(['pid', 'name', 'cpu_percent'])
    top_processes = sorted(processes, key=lambda p: p.info['cpu_percent'], reverse=True)[:15]
    return top_processes

完整示例

import psutil
import time
import curses
from pynvml import *

def get_top_cpu_processes():
    processes = psutil.process_iter(['pid', 'name', 'cpu_percent'])
    top_processes = sorted(processes, key=lambda p: p.info['cpu_percent'], reverse=True)[:15]
    return top_processes

def get_system_stats():
    cpu_usage = psutil.cpu_percent(interval=1)
    cpu_core_usages = psutil.cpu_percent(interval=1, percpu=True)
    memory_info = psutil.virtual_memory()
    memory_usage = memory_info.percent

    # 获取CPU频率
    cpu_freq = psutil.cpu_freq()

    return cpu_usage, cpu_core_usages, memory_usage, cpu_freq

def get_gpu_stats():
    nvmlInit()
    device_count = nvmlDeviceGetCount()
    gpu_stats = []
    for i in range(device_count):
        handle = nvmlDeviceGetHandleByIndex(i)
        gpu_name = nvmlDeviceGetName(handle)
        gpu_memory = nvmlDeviceGetMemoryInfo(handle)
        gpu_utilization = nvmlDeviceGetUtilizationRates(handle)
        gpu_temp = nvmlDeviceGetTemperature(handle, NVML_TEMPERATURE_GPU)
        gpu_stats.append({
            'name': gpu_name,
            'memory_used': gpu_memory.used / 1024**3,  # Convert to MB
            'memory_total': gpu_memory.total / 1024**3,  # Convert to MB
            'utilization': gpu_utilization.gpu,
            'temperature': gpu_temp
        })
    nvmlShutdown()
    return gpu_stats

def draw_bar(stdscr, y, width, percent, color_pair, color_inv):
    bar_length = int(width * (percent / 100))
    bar = ' ' * bar_length
    stdscr.addstr(y, 0, bar, curses.color_pair(color_pair) | curses.A_BOLD)
    stdscr.addstr(y, bar_length, '|' * (width - bar_length), curses.color_pair(6))
    stdscr.addstr(y, width, '|', curses.color_pair(color_inv))

def draw_rect(stdscr, y: int, x: int, width: int, max_width: int, color_pair, color_inv):
    bar = ' '*width
    stdscr.addstr(y, x, bar, curses.color_pair(color_pair) | curses.A_BOLD)
    stdscr.addstr(y, width+x, '|' * (max_width - width), curses.color_pair(6))
    stdscr.addstr(y, max_width+x, '|', curses.color_pair(color_inv))

def display_stats(stdscr):
    curses.curs_set(0)  # Hide cursor
    stdscr.nodelay(1)   # Non-blocking input

    # Initialize colors
    curses.start_color()
    curses.init_pair(1, curses.COLOR_WHITE, curses.COLOR_BLACK)
    curses.init_pair(2, curses.COLOR_RED, curses.COLOR_BLACK)
    curses.init_pair(3, curses.COLOR_WHITE, curses.COLOR_BLACK)
    curses.init_pair(4, curses.COLOR_GREEN, curses.COLOR_BLACK)
    curses.init_pair(5, curses.COLOR_BLACK, curses.COLOR_WHITE)
    curses.init_pair(6, curses.COLOR_BLACK, curses.COLOR_BLACK)
    curses.init_pair(7, curses.COLOR_BLACK, curses.COLOR_YELLOW)
    curses.init_pair(8, curses.COLOR_BLACK, curses.COLOR_GREEN)
    curses.init_pair(9, curses.COLOR_YELLOW, curses.COLOR_BLACK)
    curses.init_pair(10, curses.COLOR_GREEN, curses.COLOR_BLACK)
    curses.init_pair(11, curses.COLOR_BLACK, curses.COLOR_CYAN)
    curses.init_pair(12, curses.COLOR_CYAN, curses.COLOR_BLACK)
    
    bold = curses.A_BOLD

    
    # f = open("monLog.txt", "w")
    while True:
        height, width = stdscr.getmaxyx()
        max_width = width-10
        cpu_usage, cpu_core_usages, memory_usage, cpu_freq = get_system_stats()
        gpu_stats = get_gpu_stats()

        stdscr.clear()
        stdscr.addstr(0, int((max_width-22)/2), "System Monitoring Tool", curses.color_pair(5) | bold)
        
        stdscr.addstr(1, 0, "■ CPU Usage:", curses.color_pair(1) | bold)
        stdscr.addstr(1, int(max_width/2), "%d"%cpu_usage + "%", curses.color_pair(1) | bold)
        # f.write(str(int((int(max_width/2)-12)*cpu_usage/100))+"\n")
        # f.write("cpu usage: %d\n" % cpu_usage)
        # f.flush()
        draw_rect(stdscr, 1, 12, int((int(max_width/2)-13)*cpu_usage/100), int(int(max_width/2)-13), 11, 12)
        # draw_bar(stdscr, 3, max_width, cpu_usage, 7, 9)
        
        
        line = 3
        id = 0

        while id < len(cpu_core_usages):
            stdscr.addstr(line, 0, f"{cpu_core_usages[id]:.0f}", curses.color_pair(1))
            draw_rect(stdscr, line, 3, int((max_width-12)/4*cpu_core_usages[id]/100), int((max_width-12)/4), 11, 12)
            id+=1
            stdscr.addstr(line, int((max_width-12)/4)+4, f"{cpu_core_usages[id]:.0f}", curses.color_pair(1))
            draw_rect(stdscr, line, int((max_width-12)/4)+7, int((max_width-12)/4*cpu_core_usages[id]/100), int((max_width-12)/4), 11, 12)
            id+=1
            stdscr.addstr(line, int((max_width-12)/4)*2+8, f"{cpu_core_usages[id]:.0f}", curses.color_pair(1))
            draw_rect(stdscr, line, int((max_width-12)/4)*2+11, int((max_width-12)/4*cpu_core_usages[id]/100), int((max_width-12)/4), 11, 12)
            id+=1
            stdscr.addstr(line, int((max_width-12)/4)*3+12, f"{cpu_core_usages[id]:.0f}", curses.color_pair(1))
            draw_rect(stdscr, line, int((max_width-12)/4)*3+15, int((max_width-12)/4*cpu_core_usages[id]/100), int((max_width-12)/4), 11, 12)
            id+=1
            line+=1
        # line += 1
        stdscr.addstr(line, 0, "■ Memory Usage:", curses.color_pair(3) | bold)
        stdscr.addstr(line, int(max_width/2), f"{memory_usage:.2f}%", curses.color_pair(3) | bold)
        line+=1
        draw_bar(stdscr, line, max_width, memory_usage, 7, 9)
        line+=1
        stdscr.addstr(line, 0, "-"*(max_width), curses.color_pair(1) | bold)
        line+=1
        for i, gpu in enumerate(gpu_stats):
            stdscr.addstr(line, 0, f"■ GPU {i} ({gpu['name']}):", curses.color_pair(2) | bold)
            
            stdscr.addstr(line + 1, 0, "  Memory Usage:", curses.color_pair(4))
            stdscr.addstr(line + 1, int(max_width/2), f"{gpu['memory_used']:.2f} GB / {gpu['memory_total']:.2f} GB", curses.color_pair(4))
            draw_bar(stdscr, line + 2, max_width, (gpu['memory_used'] / gpu['memory_total']) * 100, 8, 10)
            
            stdscr.addstr(line + 3, 0, "  Utilization:", curses.color_pair(4))
            stdscr.addstr(line + 3, int(max_width/2), f"{gpu['utilization']}%", curses.color_pair(4))
            draw_bar(stdscr, line + 4, max_width, gpu['utilization'], 8, 10)
            
            stdscr.addstr(line + 5, 0, f"  Temperature: {gpu['temperature']}°C", curses.color_pair(4))
            line += 7

        # 获取CPU占用排名前五的进程信息
        top_processes = get_top_cpu_processes()
        # 跳过第一个system Idle process
        top_processes = top_processes[1:]
        # 显示CPU占用排名前五的进程信息
        line -=1
         # 确定进程名称和PID的最大长度
        max_name_length = max(max(len(p.info['name']) for p in top_processes), 20)
        max_pid_length = len(str(max(p.info['pid'] for p in top_processes)))+10

        # 显示CPU占用排名前五的进程信息
        for i, process in enumerate(top_processes):
            if line >= height:
                break
            # 格式化字符串以确保对齐
            process_info = f"{process.info['name'].ljust(max_name_length)} PID: {process.info['pid']:<{max_pid_length}} - CPU: {process.info['cpu_percent']}%"
            if line < height:  # 确保不会超出屏幕边界
                stdscr.addstr(line, 0, process_info, curses.color_pair(1) | bold)
            line += 1


        stdscr.refresh()

        # Wait for 1 second before updating the stats
        time.sleep(0.5)

        # Check if user pressed 'q' to quit
        if stdscr.getch() == ord('q'):
            break

if __name__ == "__main__":
    curses.wrapper(display_stats)

效果