A simple decorator written as a class, which counts how many times a function has been called

Posted on

Problem

Can I improve its typing? Is there any other improvement or pythonic change that you would do?

F = TypeVar('F', bound=Callable[..., Any])


# This is mostly so that I practice using a class as a decorator.
class CountCalls:
    """Logs to DEBUG how many times a function gets called, saves the result in a newly created attribute `num_calls`."""
    def __init__(self, func: F) -> None:
        functools.update_wrapper(self, func)
        self.func = func
        self.num_calls: int = 0
        self._logger = logging.getLogger(__name__ + '.' + self.func.__name__)
        self.last_return_value = None

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        self.num_calls += 1
        self._logger.debug(' called %s times', self.num_calls)
        self.last_return_value = self.func(*args, **kwargs)
        return self.last_return_value

Here’s the decorator in action:

>>> @CountCalls
... def asdf(var: str):
...     print(var)
...     return len(var)
... 
>>> asdf('Laur')
Laur
4
DEBUG:__main__.asdf: called 1 times
>>> asdf('python 3')
DEBUG:__main__.asdf: called 2 times
python 3
8
>>> asdf(3)
DEBUG:__main__.asdf: called 3 times
3
Traceback (most recent call last):
  File "<input>", line 1, in <module>
  File "C:/Projects/Python/he/src/he/decorators.py", line 156, in __call__
    self.last_return_value = self.func(*args, **kwargs)
  File "<input>", line 4, in asdf
TypeError: object of type 'int' has no len()
>>> asdf.num_calls
3

Solution

One thing you could try to improve the typing would be to type the method itself (although I’m not sure how well tools support it). Also, leading/trailing whitespace should be up to the logger, not the code using it.

F = TypeVar('F', bound=Callable[..., Any])


# This is mostly so that I practice using a class as a decorator.
class CountCalls:
    """Logs to DEBUG how many times a function gets called, saves the result in a newly created attribute `num_calls`."""
    def __init__(self, func: F) -> None:
        functools.update_wrapper(self, func)
        self.func = func
        self.num_calls: int = 0
        self._logger = logging.getLogger(__name__ + '.' + self.func.__name__)
        self.last_return_value = None

    __call__: F

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        self.num_calls += 1
        self._logger.debug(f'called %s times', self.num_calls)
        self.last_return_value = self.func(*args, **kwargs)
        return self.last_return_value

As for the code itself, you could make a callback-based API.

F = TypeVar('F', bound=Callable[..., Any])


# This is mostly so that I practice using a class as a decorator.
class CountCalls:
    """Logs to DEBUG how many times a function gets called, saves the result in a newly created attribute `num_calls`."""
    def __init__(self, func: F, callback: Optional[Callable[[int, Tuple[Any], Dict[str, Any]], Any]] = None) -> None:
        if callback is None:
            logger = logging.getLogger(__name__ + '.' + self.func.__name__)

            def callback(num_calls: int, args: Tuple[Any], kwargs: Dict[str, Any]):
                self._logger.debug(f'called %s times', self.num_calls)

        functools.update_wrapper(self, func)
        self.func = func        
        self.callback = callback
        self.num_calls: int = 0
        self.last_return_value = None

    __call__: F

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        self.num_calls += 1
        self.callback(self.num_calls, args, kwargs)
        self.last_return_value = self.func(*args, **kwargs)
        return self.last_return_value

Or with the number of calls tracked in the callback (for increased flexibility):

F = TypeVar('F', bound=Callable[..., Any])


# This is mostly so that I practice using a class as a decorator.
class CountCalls:
    """Logs to DEBUG how many times a function gets called, saves the result in a newly created attribute `num_calls`."""
    def __init__(self, func: F, callback: Optional[Callable[[int, Tuple[Any], Dict[str, Any]], Any]] = None) -> None:
        if callback is None:
            logger = logging.getLogger(__name__ + '.' + self.func.__name__)
            num_calls: int = 0

            def callback(args: Tuple[Any], kwargs: Dict[str, Any]):
                nonlocal num_calls  # Not sure if this is necessary or not
                num_calls += 1
                self._logger.debug(f'called %s times', self.num_calls)

        functools.update_wrapper(self, func)
        self.func = func        
        self.callback = callback
        self.last_return_value = None

    __call__: F

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        self.callback(self.num_calls, args, kwargs)
        self.last_return_value = self.func(*args, **kwargs)
        return self.last_return_value

Earlier, to pass a keyword argument while using it, @functools.partial(CountCalls, callback=callback) was needed. Now, @CountCalls(callback=callback) can be used instead.

F = TypeVar('F', bound=Callable[..., Any])


# This is mostly so that I practice using a class as a decorator.
class CountCalls:
    """Logs to DEBUG how many times a function gets called, saves the result in a newly created attribute `num_calls`."""
    def __init__(self, func: F = None, callback: Optional[Callable[[int, Tuple[Any], Dict[str, Any]], Any]] = None) -> None:
        if callback is None:
            logger = logging.getLogger(__name__ + '.' + self.func.__name__)
            num_calls: int = 0

            def callback(args: Tuple[Any], kwargs: Dict[str, Any]):
                nonlocal num_calls  # Not sure if this is necessary or not
                num_calls += 1
                self._logger.debug(f'called %s times', self.num_calls)

        if func is None:
            return functools.partial(CountCalls, callback=callback)

        functools.update_wrapper(self, func)
        self.func = func        
        self.callback = callback
        self.last_return_value = None

    __call__: F

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        self.callback(self.num_calls, args, kwargs)
        self.last_return_value = self.func(*args, **kwargs)
        return self.last_return_value

(Note: none of this code has been tested.)

Leave a Reply

Your email address will not be published. Required fields are marked *