timer.py

"""
Module containing utility functions to time code snippets.
"""

import contextlib
import time


@contextlib.contextmanager
def timelogger(msg_prefix:str='Executed in : ', msg_suffix:str='s', logger=print):
    """
    Context manager for measuring the execution time for any code block.

    Decorators are a convenient way to time functions but to time only few
    lines inside the function body, this context manager can be used.
    It will time the code inside the `with` block using this context manager.

    It will log the formatted message immediately after `with` block ends.

    This context manager provides following flexibility:
    - display message can be easily customized.
    - any logger function can be passed.

    :param str msg_prefix: Text to be displayed left of execution time.
    :param str msg_suffix: Text to be displayed right of execution time.
    :param logger: Pointer to function with which you want to log the
        formatted message immediately after the context block. e.g.: print

    >>> with timelogger('Squares computed in : '): # doctest: +ELLIPSIS
    ...     count = 100000
    ...     sq = [x*x for x in range(count)]
    Squares computed in : ...s
    >>> with timelogger(msg_suffix='sec', logger=print): # doctest: +ELLIPSIS
    ...     count = 100000
    ...     sq = [x*x*x for x in range(count)]
    Executed in : ...sec
    >>> with timelogger(): # doctest: +ELLIPSIS
    ...     res = 5 + '4'
    Traceback (most recent call last):
    ...
        res = 5 + '4'
    TypeError: unsupported operand type(s) for +: 'int' and 'str'
    """
    start_time = time.perf_counter()
    try:
        yield
    finally:
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        logger('{}{}{}'.format(msg_prefix, execution_time, msg_suffix))


class TimeLogger:
    """
    Context manager for measuring the execution time for any code block.

    Decorators are a convenient way to time functions but to time only few
    lines inside the function body, this context manager can be used.
    It will time the code inside the `with` block using this context manager.

    Unless logger is set to None, it will log the formatted message
    immediately after `with` block ends.

    This class provides following flexibility:
    - display message can be easily customized.
    - any logger function can be passed.
    - the context manager instance can be used to access the attributes later.
    - can control whether the logging should happen or not immediately after `with` block.

    >>> with TimeLogger('Squares computed in : '): # doctest: +ELLIPSIS
    ...     count = 100000
    ...     sq = [x*x for x in range(count)]
    Squares computed in : ...s
    >>> with TimeLogger(logger=None) as tl:
    ...     count = 100000
    ...     sq = [x*x*x for x in range(count)]
    >>> print(tl.get_msg()) # doctest: +ELLIPSIS
    Executed in : ...s
    >>> print('Time taken :', tl.execution_time, 'sec') # doctest: +ELLIPSIS
    Time taken : ... sec
    >>> with TimeLogger(): # doctest: +ELLIPSIS
    ...     res = 5 + '4'
    Traceback (most recent call last):
    ...
        res = 5 + '4'
    TypeError: unsupported operand type(s) for +: 'int' and 'str'
    """

    def __init__(self, msg_prefix:str='Executed in : ', msg_suffix:str='s', logger=print):
        """
        :param str msg_prefix: Text to be displayed left of execution time.
        :param str msg_suffix: Text to be displayed right of execution time.
        :param logger: Pointer to function with which you want to log the
            formatted message immediately after the context block. e.g.: print
            If None, won't log anything after the function (this is handy when
            you want to time the snippet but log it at a later point using
            the context variable)
        """
        self._msg_prefix = msg_prefix
        self._msg_suffix = msg_suffix
        self._logger = logger

    def __enter__(self):
        """
        Function which will run at end of context block.

        Captures the time which will be used to compute the execution time.
        """
        self._start_time = time.perf_counter()
        return self

    def __exit__(self, *args):
        """
        Function which will run at end of context block.

        Computes the time spent in the context block and logs the formatted
        message unless logger is set to None.
        """
        self._end_time = time.perf_counter()
        self.execution_time = self._end_time - self._start_time
        if self._logger is not None:
            self._logger(self.get_msg())

    def get_msg(self) -> str:
        """
        Returns the formatted string to be logged.

        Concatenates the following: msg_prefix, execution_time, msg_suffix
        """
        return '{}{}{}'.format(self._msg_prefix, self.execution_time, self._msg_suffix)