from twisted.internet.defer import maybeDeferred

def withDeferred(contextManager, f, *args, **kwargs):
    contextManager.__enter__()

    def _exit(result):
        contextManager.__exit__(None, None, None)
        return result

    def _exitEb(failure):
        if contextManager.__exit__(
            failure.type, failure.value, failure.stack):
            return failure

    d = maybeDeferred(f, *args, **kwargs)
    d.addCallbacks(_exit, _exitEb)

    return d

if __name__ == '__main__':
    from greplin import scales
    from twisted.internet.defer import Deferred
    from twisted.internet import reactor

    STATS = scales.collection('/my-scales', scales.PmfStat('timing'))

    def doStuff(_, arg):
        print "bar"
        d = Deferred()
        if arg < 3:
            d.addCallback(doStuff, arg + 1)

        reactor.callLater(1, d.callback, True)
        return d

    def done(_):
        print scales.getStats()

    def doit():
        print "foo"
        d2 = withDeferred(STATS.timing.time(), doStuff, None, 0)
        d2.addCallback(done)

    reactor.callLater(0, doit)
    reactor.callLater(2, doit)
    reactor.callLater(30, doit)
    reactor.callLater(20, doit)
    reactor.run()