profiler.py

import time


global_hooks = {}


class FunctionHook(object):

    def __enter__(self):
        global global_hooks
        global_hooks[str(self)] = self
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        global global_hooks
        del global_hooks[str(self)]

    def __call__(self, function):
        raise NotImplementedError


class CallableHook(FunctionHook):

    def __init__(self, f):
        self.f = f

    def __call__(self, function):
        return self.f(function)


class PrintHook(FunctionHook):

    def __call__(self, function):
        print('from hook object', function.label)


class TimerHook(FunctionHook):

    def preprocess(self, function):
        self.start = time.time()

    def __call__(self, function):
        self.end = time.time()
        print('{}\t{}'.format(function.label, self.end - self.start))


class Function(object):

    def __init__(self):
        self.local_hook = {}

    @property
    def label(self):
        return 'function'

    def add_hook(self, hook, name=None):
        if name is None:
            name = str(hook)
        self.local_hook[name] = hook

    def delete_hook(self, name):
        if name in self.local_hook:
            del self.local_hook[name]

    def forward(self):
        print('forward\t{}'.format(self.label))

    def __call__(self):

        hooks = global_hooks.values() + self.local_hook.values()

        for hook in hooks:
            if hasattr(hook, 'preprocess'):
                hook.preprocess(self)
        self.forward()
        for hook in hooks:
            hook(self)


class F(Function):

    @property
    def label(self):
        return 'F'


class G(Function):

    @property
    def label(self):
        return 'G'


def print_(function):
    print('from callable', function.label)


f = F()
g = G()
f.add_hook(TimerHook())

with CallableHook(print_), PrintHook():
    f()
    g()

'''
forward	F
('from hook object', 'F')
('from callable', 'F')
F	5.00679016113e-05
forward	G
('from hook object', 'G')
('from callable', 'G')
'''