Python自定义装饰器-统计函数的耗时和CPU使用率

经常遇到耗时分析、资源利用率分析等，最近手搓了一份统计函数的耗时和CPU使用率的Python装饰器，纯个人的自定义用法，后续再改进。

import time
import psutil
import numpy as np


def profile(func):
    def wrapper(*args, **kwargs):
        result = None
        elapsed, cpu_percentage, cpu_nozero_times = 0.0, 0.0, 0
        psp = psutil.Process(os.getpid())
        t0 = time.time()
        psp.cpu_percent()  # last call
        result = func(*args, **kwargs)
        cpu_percent = psp.cpu_percent()  # percentage since last call
        t1 = time.time()  # 略微牺牲下耗时
        elapsed += t1 - t0
        if cpu_percent > 0:
            cpu_percentage += cpu_percent
            cpu_nozero_times += 1
        num_cpus_system = psutil.cpu_count()
        if cpu_nozero_times > 0:
            cpu_percentage /= cpu_nozero_times * num_cpus_system
        else:
            cpu_percentage = -1
        args_str = ''
        if len(args) > 0:
            args_str += ', '.join([f"ndarray{x.shape}+{str(x.dtype)}" if isinstance(x, np.ndarray) else str(x) for x in args])
        if len(kwargs.keys()) > 0:
            kwargs_str = ', '.join([f"{k}=ndarray{v.shape}+{str(v.dtype)}" if isinstance(v, np.ndarray) else f"{k}={v}" for k, v in kwargs.items()])
            args_str += f', {kwargs_str}' if len(args) > 0 else kwargs_str
        print(f'{func.__name__}({args_str}), 跑了1次, '
              f'平均耗时: {round(elapsed * 1000, 4)} ms, '
              f'CPU使用率: {round(cpu_percentage, 2)}% ({cpu_nozero_times} times on {num_cpus_system} CPU)')
        return result
    return wrapper

用法很简单：

@profile
def foo(*args, **kwargs):
    ...