Simple python profiler

profiler.py

import functools
from time import time
from .text_styles import *


class ProfileItem(object):

    def __init__(self, name, elapsed, stats):
        self.name = name
        self.elapsed = elapsed
        self.stats = stats

    def __str__(self):
        sorted_stats = sorted(self.stats.items(), key=lambda t: t[0])
        stats_string = "\n  ".join(map(lambda kv: "{}: {}".format(kv[0].title(), kv[1]), sorted_stats))
        if stats_string != "":
            stats_string = "\n  " + stats_string
        return "{} {} seconds{}".format(cyan(bold(self.name + ":")),
                                        self.elapsed,
                                        stats_string)


class Profiler(object):

    def __init__(self, defer_printing=False):
        self.defer_printing = defer_printing
        self.print_queue = []

    def _handle_print(self, print_item):
        line = str(print_item)
        if self.defer_printing:
            print(green(">>> Finished {}".format(print_item.name)))
            self.print_queue.append(line)
        else:
            print(line)

    def flush_print_queue(self):
        if self.defer_printing:
            print(underlined(magenta("\nProfile")))
            for line in self.print_queue:
                print(line)
            print()
            self.print_queue = []

    def profile(self, name):
        def wrapper(func):
            @functools.wraps(func)
            def wrapped(*args, **kwargs):
                ret = None
                start_time = time()
                profiler_stats = {}
                try:
                    ret = func(*args, profiler_stats=profiler_stats, **kwargs)
                except TypeError as e:
                    if "profiler_stats" in str(e):
                        ret = func(*args, **kwargs)
                    else:
                        raise
                end_time = time()
                elapsed = end_time - start_time

                profile_item = ProfileItem(name, elapsed, profiler_stats)
                self._handle_print(profile_item)

                return ret
            return wrapped
        return wrapper

Requires text_styles.py