pythonprofilers / memory_profiler

Monitor Memory usage of Python code
http://pypi.python.org/pypi/memory_profiler
Other
4.39k stars 380 forks source link

The @profile decorator will profile another decorator when chaining multiple decorators into a function. #373

Open stenbein opened 2 years ago

stenbein commented 2 years ago

Stacking multiple decorators together will result in unexpected behavior. A trivial example:

from functools import wraps
from random import random
import time

from memory_profiler import profile

def timeit(f):
    @wraps(wrapped=f)
    def wrapper(*args, **kw):
        ts = time.time()
        result = f(*args, **kw)
        te = time.time()
        print(f'func: {f.__name__} took: {te-ts:2.4f} sec')
        return result

    return wrapper

@profile
@timeit
def to_profile(n):
    arr = [random() for i in range(n)]
    return sum(arr)

to_profile(1_000)

What I expect to see:

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    20     48.8 MiB     48.8 MiB           1   @timeit
    21                                         @profile
    22                                         def to_profile(n):
    23     48.8 MiB      0.0 MiB        1003       arr = [random() for i in range(n)]
    24     48.8 MiB      0.0 MiB           1       return sum(arr)

What we actually get:

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     9     49.0 MiB     49.0 MiB           1       @wraps(wrapped=f)
    10                                             def wrapper(*args, **kw):
    11     49.0 MiB      0.0 MiB           1           ts = time.time()
    12     49.0 MiB      0.0 MiB           1           result = f(*args, **kw)
    13     49.0 MiB      0.0 MiB           1           te = time.time()
    14     49.0 MiB      0.0 MiB           1           print(f'func: {f.__name__} took: {te-ts:2.4f} sec')
    15     49.0 MiB      0.0 MiB           1           return result

I realized while checking this that the sample timeit decorator has the same issue and is timing the @profile rather than the underlying to_profile.

stenbein commented 2 years ago

Possible partial solution: Decorators which use functools.wraps annotate the function object with wrapped property which points to the wrapped target. Recursively checking for the wrapper and replacing the function in each case correctly profiles the target function. See mock https://github.com/stenbein/memory_profiler/tree/stacked-decorators