Change an attribute of a function inside its own body?

Question

I'm attempting to create a function that keeps count of the times it has been called, and I want the information to stay inside the function itself. I tried creating a wrapper like so:

def keep_count(f):
    f.count = 0
    @functools.wraps(f)
    def wrapped_f(*args, **kwargs):
        f(*args, **kwargs)
        f.count += 1
    return wrapped_f

@keep_count
def test_f(*args, **kwargs):
    print(args, kwargs)

I thought it'd work, but I got an AttributeError saying 'function' object has no attribute 'count'.

I already figured the problem: It's because the decorator sets my test_f to equal to the wrapped_f (defined inside the keep_count decorator), and I'm increasing the count of the original f, which is no longer used since test_f refers to the new function.

Is there a way to do this without too much of hacking?

Answer 1

Just set the attribute on wrapped_f instead of f. This requires you to set the initial count after defining the function, but that's no big deal:

def keep_count(f):
    @functools.wraps(f)
    def wrapped_f(*args, **kwargs):
        f(*args, **kwargs)
        wrapped_f.count += 1
    wrapped_f.count = 0
    return wrapped_f

Then with your decorated function:

>>> test_f()
() {}
>>> test_f.count
1
>>> test_f()
() {}
>>> test_f.count
2

This works because wrapped_f is a local variable inside keep_count. Since the body of wrapped_f contains a reference to wrapped_f, it gets a closure allowing wrapped_f to access itself.