« Back to Index

[Python Class Implementation of a Decorator]

View original Gist on GitHub

Tags: #python #decorator #class

Python Class Implementation of a Decorator.py

"""
[counter 1] add called 1 times
[counter 1] add called 2 times
[counter2] add2 called 1 times
[counter2] add2 called 2 times
[counter2] add2 called 3 times
"""

from functools import update_wrapper


class call_count:
    def __init__(self, name):
        self.__name = name
        self.__count = 0

    def __call__(self, fn, *args, **kwargs):
        def wrapper(*args, **kwargs):
            self.__count += 1
            result = fn(*args, **kwargs)
            print(f'[{self.__name}] {fn.__name__} called {self.__count} times')
            return result
        update_wrapper(self, wrapper)
        return wrapper

    @property
    def call_count(self):
        return self.__count


@call_count('counter 1')
def add(a, b):
    return a + b


if __name__ == '__main__':
    assert add(1, 2) == 3
    assert add(3, 4) == 7

    counter2 = call_count('counter2')

    @counter2
    def add2(a, b):
        return a + b

    add2(1, 1), add2(1, 1), add2(1, 1)
    assert counter2.call_count == 3