optimize-later/optimize_later/config.py

import logging
from functools import wraps

from optimize_later.utils import NoArgDecoratorMeta, with_metaclass

try:
    import threading
except ImportError:
    import dummy_threading as threading

log = logging.getLogger(__name__.rpartition('.')[0] or __name__)

_global_callbacks = []
_local = threading.local()


def get_callbacks():
    try:
        return _local.callbacks
    except AttributeError:
        return _global_callbacks


def register_callback(callback):
    get_callbacks().append(callback)
    return callback


def deregister_callback(callback):
    get_callbacks().remove(callback)


def global_callback(report):
    for callback in get_callbacks():
        try:
            callback(report)
        except Exception:
            log.exception('Failed to invoke global callback: %r', callback)


class optimize_context(with_metaclass(NoArgDecoratorMeta)):
    def __init__(self, callbacks=None, reset=False):
        self.callbacks = callbacks or []
        self.reset = reset

    def __enter__(self):
        try:
            self.old_context = _local.callbacks
        except AttributeError:
            self.old_context = None

        if self.reset:
            base_context = []
        elif self.old_context is None:
            base_context = _global_callbacks
        else:
            base_context = self.old_context
        _local.callbacks = base_context + self.callbacks

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.old_context is None:
            del _local.callbacks
        else:
            _local.callbacks = self.old_context

    def __call__(self, function):
        @wraps(function)
        def wrapper(*args, **kwargs):
            with optimize_context(self.callbacks, reset=self.reset):
                return function(*args, **kwargs)
        return wrapper