From e63764c72e4223f86cf6e60ab0bfbdd734ee740e Mon Sep 17 00:00:00 2001 From: Quantum Date: Fri, 11 Aug 2017 16:08:59 -0400 Subject: [PATCH] Initial version. --- optimize_later/__init__.py | 0 optimize_later/config.py | 10 ++ optimize_later/core.py | 184 +++++++++++++++++++++++++++++++++++++ optimize_later/tests.py | 148 +++++++++++++++++++++++++++++ 4 files changed, 342 insertions(+) create mode 100644 optimize_later/__init__.py create mode 100644 optimize_later/config.py create mode 100644 optimize_later/core.py create mode 100644 optimize_later/tests.py diff --git a/optimize_later/__init__.py b/optimize_later/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/optimize_later/config.py b/optimize_later/config.py new file mode 100644 index 0000000..32d1720 --- /dev/null +++ b/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) diff --git a/optimize_later/core.py b/optimize_later/core.py new file mode 100644 index 0000000..3d31131 --- /dev/null +++ b/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 diff --git a/optimize_later/tests.py b/optimize_later/tests.py new file mode 100644 index 0000000..c2668d3 --- /dev/null +++ b/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])