How to reference an object instance in callable decorator object with python?

Question

Context:

I'd like to be able to decorate functions so that I can track their stats. Using this post as a reference I went about trying to make my own callable decorator objects.

Here is what I ended up with:

def Stats(fn):
    Class StatsObject(object):
        def __init__(self, fn):
            self.fn = fn
            self.stats = {}

        def __call__(self, obj, *args, **kwargs):
            self.stats['times_called'] = self.stats.get('times_called', 0) + 1
            return self.fn(obj, *args, **kwargs)

    function = StatsObject(fn)
    def wrapper(self, *args **kwargs):
        return function(self, *args, **kwargs)
    return wrapper

Class MockClass(object):
    @Stats
    def mock_fn(self, *args, **kwargs):
        # do things

Problem:

This actually calls the mock_fn function correctly but I don't have a reference to the stats object outside the wrapper function. i.e. I can't do:

mc = MockClass()
mc.mock_fn()
mc.mock_fn.stats
# HasNoAttribute Exception

Then I tried changing the following code recognizing that it was a scoping issue:

From:

    function = StatsObject(fn)
    def wrapper(self, *args **kwargs):
        return function(self, *args, **kwargs)
    return wrapper

To:

    function = StatsObject(fn)
    return function

But of course I lost the self reference (self becomes the StatsObject instance, obj becomes the first arg, and the MockClass object self reference gets lost).

So I understand why the first issue is happening, but not the second. Is there any way that I can pass the self reference of MockClass to the StatsObject __call__ function?

Answer 1

Functions can actually themselves have attributes in Python.

def Stats(fn):
    class StatsObject(object):
        def __init__(self, fn):
            self.fn = fn
            self.stats = {}

        def __call__(self, obj, *args, **kwargs):
            self.stats['times_called'] = self.stats.get('times_called', 0) + 1
            return self.fn(obj, *args, **kwargs)

    function = StatsObject(fn)
    def wrapper(self, *args **kwargs):
        return function(self, *args, **kwargs)

    # KEY LINE BELOW: make the StatsObject available outside as "stats_fn"
    wrapper.stats_fn = function

    return wrapper

class MockClass(object):
    @Stats
    def mock_fn(self, *args, **kwargs):
        # do things

The key line is assigning the StatsObject instance (which you've, perhaps misleadingly, locally named function) as an attribute of the function which you return from the decorator.

Once you do this, self.mock_fn.stats_fn.stats (not self.mock_fn()! The attribute is on the function, not its return value) will work within an instance of MockClass, and MockClass.mock_fn.stats_fn.stats will be available outside. The statistics will be global across all instances of MockClass (since the decorator is called once, not once per instance), which may or may not be what you want.