Python：向装饰器添加方法

Question

我编写了一个装饰器，用于对给定函数的所有调用执行分析（cProfile）。 .print_statistics() 方法的调用应输出一个表，其中包含汇总该函数的所有调用的分析统计信息。
可以这样添加装饰器方法：

def add_method(fn):
    # first method:
    fn.__dict__.update({"print_stats_1": lambda: print("print_statistics_1 works!")})
    # second method:
    fn.print_stats_2 = lambda: print("print_statistics_2 works!")
    print(fn.__dict__)
    return fn

@add_method
def greet(name):
    print(f"Hi there, {name}!")

# Test the decorated function
greet("John")

greet.print_statistics_1()
greet.print_statistics_2()

我们得到这个结果：

Hi there, John!
print_statistics_1 works!
print_statistics_2 works!

一切正常，

print(fn.__dict__)

还显示了属性字典中这些方法的存在。
但是当我将这个逻辑实现到我的分析装饰器中时，我收到一个错误。这是我的代码：

import cProfile
import io
import pstats
from functools import wraps


def profiled(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        pr = cProfile.Profile()
        pr.enable()
        result = func(*args, **kwargs)
        pr.disable()
        s = io.StringIO()
        ps = pstats.Stats(pr, stream=s).sort_stats("cumulative")
        ps.print_stats()
        # print(s.getvalue()) # this works, but I need print stats using a method, not directly:
        func.print_statistics = lambda: print(s.getvalue()) # 1st method
        # func.__dict__.update({"also_print_statistics": lambda: print(s.getvalue())})  # 2d method
        print(func.__dict__)
        return result
    return wrapper



@profiled
def add(a, b):
    return a + b

add(1, 2)
add.print_statistics()

这是一个错误：

Traceback (most recent call last):
  File "C:\Users\user\PycharmProject\profiler.py", line 49, in <module>
    add.print_statistics()
    ^^^^^^^^^^^^^^^^^^^^
AttributeError: 'function' object has no attribute 'print_statistics'

即使属性 dict

func.__dict__

包含此方法，我也会收到错误。为什么？我怎样才能使这个方法发挥作用？

Answer 1

您可以执行以下操作来修改“func”

将语句移到“包装器”之前，如下所示：

import cProfile
import io
import pstats
from functools import wraps


def profiled(func):
    pr = cProfile.Profile()
    pr.enable()
    pr.disable()
    s = io.StringIO()
    ps = pstats.Stats(pr, stream=s).sort_stats("cumulative")
    ps.print_stats()
    # print(s.getvalue()) # this works, but I need print stats using a method, not directly:
    func.print_statistics = lambda: print(s.getvalue()) # 1st method
    # func.__dict__.update({"also_print_statistics": lambda: print(s.getvalue())})  # 2d method
    print(func.__dict__)
    @wraps(func)
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        return result

    return wrapper



@profiled
def add(a, b):
    return a + b

add(1, 2)
add.print_statistics()

打印以下内容

{'print_statistics': <function profiled.<locals>.<lambda> at 0x000001F85DBD40E0>}
         1 function calls in 0.000 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}