Timer.py: Quick ways to time blocks of code.

timer.py

from __future__ import print_function
import time
import sys

__author__ = "Henry Lin <hlin117@gmail.com>"

class Timer(object):
    """Implements a timer which clocks the amount of time it takes for
    a piece of code to run. It wraps a "context", which makes timing
    the code unobtrusive.

    This Timer also addresses the bug that ipython notebooks have;
    ipython notebooks have a "fixed buffering" problem, where output
    is not sent to STDOUT immediately.

    Another open source implementation is offered here. However, I don't
    guarantee that it has the same features as this timer.
    https://pypi.python.org/pypi/contexttimer/0.3.1

    You're free to modify this code.

    Parameters
    ----------
    annotation : str
        Assigns an annotation to the context block that is being timed.

    color : bool
        Indicates whether color should be associated with the output.
        At the current moment, the timers will cycle through the following
        colors:

            Red, Green, Blue, Magenta, Cyan

        You must have the colorama package installed to use this feature.

    Examples
    --------
    The timer works with context blocks:

    >>> import sys
    >>> with Timer(annotation="Sleeper"):
    >>>     sys.sleep(5)

    The printed format will be as followed:

        > Starting Timer 1: Sleeper

        (5 seconds later)

        > Took 0 minutes and 5 seconds.

    You could even nest timers together, and with color=True, you'll
    be able to find the pairs of messages quickly.

    """

    # Incremented after each use of the timer
    _id = 1

    def __init__(self, annotation=None, color=False):
        self.annotation = annotation
        self.color = color
        self.start_ = None

        if color:
            from colorama import Fore
            colors = [Fore.RED, Fore.GREEN, Fore.BLUE, Fore.MAGENTA, Fore.CYAN]
            colorid = Timer._id % len(colors)
            self.color_ = colors[colorid]
            self.reset_ = Fore.RESET

        else:
            self.color_ = ""
            self.reset_ = ""

    def __enter__(self):
        self.start_ = time.time()
        description = ": {0}".format(self.annotation) if self.annotation else ""
        print(self.color_ + "Starting Timer {0}{1}".format(Timer._id, description + self.reset_))

        # Avoiding buffering issue:
        # http://stackoverflow.com/a/26352157/2014591
        #sys.stdout = open('/dev/stdout', 'w')
        sys.stdout.flush()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end_ = time.time()

        elapsed_time = int(self.end_ - self.start_)
        minutes = elapsed_time / 60
        seconds = elapsed_time % 60

        timestring = "Took {0} minutes and {1} seconds.".format(minutes, seconds)
        print(self.color_ + timestring + self.reset_)
        Timer._id += 1

        return (exc_type is None)