Initial version.

optimize_later/config.py

@@ -0,0 +1,10 @@
+callbacks = []
+
+
+def register_callback(callback):
+    callbacks.append(callback)
+    return callback
+
+
+def deregister_callback(callback):
+    callbacks.remove(callback)

optimize_later/core.py

@@ -0,0 +1,184 @@
+import inspect
+import logging
+import os
+import time
+from copy import copy
+from functools import wraps
+from numbers import Number
+from types import FunctionType
+
+from optimize_later import config
+
+log = logging.getLogger(__name__.rpartition('.')[0] or __name__)
+timer = [time.time, time.clock][os.name == 'nt']
+
+
+def global_callback(report):
+    for callback in config.callbacks:
+        try:
+            callback(report)
+        except Exception:
+            log.exception('Failed to invoke global callback: %r', callback)
+
+
+def _generate_default_name():
+    for entry in inspect.stack():
+        file, line = entry[1:3]
+        if file != __file__:
+            break
+    else:
+        return '-'
+    return '%s@%d' % (os.path.basename(file), line)
+
+
+class OptimizeBlock(object):
+    def __init__(self, name):
+        self.name = name
+        self.start = None
+        self.end = None
+        self.delta = None
+        self.blocks = []
+
+    def block(self, name=None):
+        block = OptimizeBlock(name or _generate_default_name())
+        self.blocks.append(block)
+        return block
+
+    def __enter__(self):
+        assert self.start is None, 'Do not reuse blocks.'
+        self.start = timer()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.end = timer()
+        self.delta = self.end - self.start
+
+    def short(self, precision=3):
+        return 'Block %r took %.*fs' % (self.name, precision, self.delta)
+
+    def long(self, precision=6):
+        lines = ['  - %s%s' % (self.short(precision), ', children:' if self.blocks else '')]
+        for block in self.blocks:
+            lines.append('    ' + block.long().replace('\n', '\n    '))
+        return '\n'.join(lines)
+
+    def __str__(self):
+        return 'Block %r took %.6fs' % (self.name, self.delta)
+
+    def __repr__(self):
+        return 'optimize_block(%r, delta=%.6f, blocks=%r)' % (self.name, self.delta, self.blocks)
+
+
+class OptimizeReport(object):
+    def __init__(self, name, limit, start, end, delta, blocks):
+        self.name = name
+        self.limit = limit
+        self.start = start
+        self.end = end
+        self.delta = delta
+        self.blocks = blocks
+
+    def short(self, precision=3):
+        return 'Block %r took %.*fs (+%.*fs over limit)' % (
+            self.name,
+            precision, self.delta,
+            precision, self.delta - self.limit,
+        )
+
+    def long(self, precision=6):
+        lines = [self.short(precision)]
+        if self.blocks:
+            lines[-1] += ', children:'
+            for block in self.blocks:
+                lines.append(block.long())
+        return '\n'.join(lines)
+
+    def __str__(self):
+        return self.short()
+
+
+class NoArgDecoratorMeta(type):
+    def __call__(cls, *args, **kwargs):
+        if len(args) == 1 and isinstance(args[0], FunctionType):
+            return cls()(args[0])
+        return super(NoArgDecoratorMeta, cls).__call__(*args, **kwargs)
+
+
+class optimize_later(object):
+    __metaclass__ = NoArgDecoratorMeta
+
+    def __init__(self, name=None, limit=None, callback=None):
+        if limit is None and isinstance(name, Number):
+            name, limit = None, name
+        self._default_name = not name
+        self.name = name or _generate_default_name()
+        self.limit = limit or 0
+        self.callback = callback
+        self.start = None
+        self.end = None
+        self.delta = None
+
+        # This is going to get shallow copied, so we shouldn't use [].
+        self.blocks = None
+
+    def block(self, name=None):
+        assert self.start is not None, 'Blocks are meant to be used inside with.'
+        if self.blocks is None:
+            self.blocks = []
+        block = OptimizeBlock(name or _generate_default_name())
+        self.blocks.append(block)
+        return block
+
+    def __enter__(self):
+        assert self.start is None, 'Do not reuse optimize_later objects.'
+        self.start = timer()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.end = timer()
+        self.delta = self.end - self.start
+        if self.delta >= self.limit:
+            self._report()
+
+    def _report(self):
+        report = OptimizeReport(self.name, self.limit, self.start, self.end, self.delta, self.blocks or [])
+        if self.callback:
+            try:
+                self.callback(report)
+            except Exception:
+                log.exception('Failed to invoke user-specified callback: %r', self.callback)
+        else:
+            global_callback(report)
+
+    def __call__(self, function):
+        self.name = '%s:%s' % (function.__module__, function.__name__)
+
+        @wraps(function)
+        def wrapped(*args, **kwargs):
+            with copy(self):
+                return function(*args, **kwargs)
+        return wrapped
+
+
+class optimize_context(object):
+    __metaclass__ = NoArgDecoratorMeta
+
+    def __init__(self, callbacks=None):
+        self.callbacks = callbacks
+
+    def __enter__(self):
+        self.old_context = config.callbacks[:]
+        if self.callbacks is None:
+            config.callbacks[:] = self.old_context
+        else:
+            config.callbacks[:] = self.callbacks
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        config.callbacks[:] = self.old_context
+
+    def __call__(self, function):
+        @wraps(function)
+        def wrapper(*args, **kwargs):
+            with optimize_context(self.callbacks):
+                return function(*args, **kwargs)
+        return wrapper

optimize_later/tests.py

@@ -0,0 +1,148 @@
+import time
+from unittest import TestCase
+
+from optimize_later import config
+from optimize_later.core import optimize_later, OptimizeReport, OptimizeBlock, optimize_context
+
+
+class OptimizeLaterTest(TestCase):
+    def setUp(self):
+        self.optimize_context = optimize_context([])
+        self.optimize_context.__enter__()
+
+    def tearDown(self):
+        self.optimize_context.__exit__(None, None, None)
+
+    def assertReport(self, report, name=None, blocks=0):
+        self.assertIsInstance(report, OptimizeReport)
+        self.assertIsInstance(report.start, float)
+        self.assertIsInstance(report.end, float)
+        self.assertIsInstance(report.delta, float)
+        self.assertIsInstance(report.blocks, list)
+
+        self.assertLessEqual(report.start, report.end)
+        self.assertGreaterEqual(report.delta, 0)
+
+        if name:
+            self.assertEqual(report.name, name)
+
+        self.assertEqual(len(report.blocks), blocks)
+
+        cumtime = 0
+        for block in report.blocks:
+            self.assertBlock(block)
+            self.assertGreaterEqual(block.start, report.start)
+            self.assertLessEqual(block.end, report.end)
+            cumtime += block.delta
+        self.assertLessEqual(cumtime, report.delta)
+
+    def assertBlock(self, block):
+        self.assertIsInstance(block, OptimizeBlock)
+        self.assertIsInstance(block.start, float)
+        self.assertIsInstance(block.end, float)
+        self.assertIsInstance(block.delta, float)
+
+        self.assertLessEqual(block.start, block.end)
+        self.assertGreaterEqual(block.delta, 0)
+
+        cumtime = 0
+        for subblock in block.blocks:
+            self.assertBlock(subblock)
+            cumtime += subblock.delta
+        self.assertLessEqual(cumtime, block.delta)
+
+    def test_default_name(self):
+        self.assertIn('.py@', optimize_later().name)
+
+    def get_report(self, *args, **kwargs):
+        reports = []
+        function = kwargs.pop('function', lambda: None)
+        with optimize_later(*args, callback=reports.append, **kwargs):
+            function()
+        self.assertIn(len(reports), (0, 1))
+        return reports[0] if reports else None
+
+    def test_simple(self):
+        self.assertReport(self.get_report())
+
+    def test_simple_fast(self):
+        self.assertIs(self.get_report(float('inf')), None)
+
+    def test_name(self):
+        self.assertReport(self.get_report('magic_name'),
+                          name='magic_name')
+
+    def test_name_fast(self):
+        self.assertIs(self.get_report('name', float('inf')), None)
+
+    def test_100_ms(self):
+        self.assertReport(self.get_report(function=lambda: time.sleep(0.1)))
+
+    def test_blocks(self):
+        reports = []
+        with optimize_later(callback=reports.append) as o:
+            with o.block('a'):
+                pass
+            with o.block('b'):
+                pass
+        self.assertEqual(len(reports), 1)
+        self.assertReport(reports[0], blocks=2)
+        for block, name in zip(reports[0].blocks, 'ab'):
+            self.assertEqual(block.name, name)
+
+    def test_block_naming(self):
+        with optimize_later() as o:
+            with o.block() as b:
+                self.assertIn('.py@', b.name)
+
+    def test_nested_block(self):
+        reports = []
+        with optimize_later(callback=reports.append) as o:
+            with o.block() as b:
+                with b.block():
+                    pass
+                with b.block():
+                    pass
+
+            with o.block():
+                pass
+
+        self.assertEqual(len(reports), 1)
+        self.assertReport(reports[0], blocks=2)
+        report = reports[0].long()
+        print report
+        self.assertIn('      - Block ', report)
+        self.assertIn('  - Block', report)
+        self.assertEqual(report.count(', children:'), 2)
+
+    def test_decorator(self):
+        reports = []
+        config.register_callback(reports.append)
+
+        @optimize_later
+        def function():
+            pass
+        self.assertEqual(function.__name__, 'function')
+        self.assertEqual(function.__module__, __name__)
+
+        for i in xrange(10):
+            function()
+        self.assertEqual(len(reports), 10)
+        for report in reports:
+            self.assertReport(report)
+
+    def test_optimize_context(self):
+        config.register_callback(1)
+        with optimize_context():
+            self.assertEqual(config.callbacks, [1])
+            config.register_callback(2)
+            self.assertEqual(config.callbacks, [1, 2])
+
+            with optimize_context([]):
+                self.assertEqual(config.callbacks, [])
+                config.register_callback(3)
+                self.assertEqual(config.callbacks, [3])
+
+            config.register_callback(4)
+            self.assertEqual(config.callbacks, [1, 2, 4])
+        self.assertEqual(config.callbacks, [1])