Python: adding method to decorator

76 Views Asked by At

I write a decorator that performs profiling (cProfile) of all calls of given function. The call of the .print_statistics() method should output a table with profiling statistics summarizing all calls of the function.
Method to decorator can be added this way:

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()

and we get this result:

Hi there, John!
print_statistics_1 works!
print_statistics_2 works!

all works fine, and print(fn.__dict__) also shows the existence of these methods in attributes dict.
But when I implement this logic into my profiling decorator, I get an error. This is my code:

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()

This is an error:

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'

I'm getting an error even though attributes dict func.__dict__ contains this method. Why? How can I make this method work?

0

There are 0 best solutions below