from collections import defaultdict
from functools import wraps
from django.db.models import signals


class disable_signals:
    def __init__(self, disabled_signals=None):
        self.stashed_signals = defaultdict(list)

        self.disabled_signals = disabled_signals or [
            signals.pre_init, signals.post_init,
            signals.pre_save, signals.post_save,
            signals.pre_delete, signals.post_delete,
            signals.pre_migrate, signals.post_migrate,
        ]

    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            self.__enter__()
            result = func(*args, **kwargs)
            self.__exit__(None, None, None)
            return result

        return wrapper

    def __enter__(self):
        for signal in self.disabled_signals:
            self.disconnect(signal)

    def __exit__(self, exc_type, exc_val, exc_tb):
        for signal in list(self.stashed_signals):
            self.reconnect(signal)

    def disconnect(self, signal):
        self.stashed_signals[signal] = signal.receivers
        signal.receivers = []

    def reconnect(self, signal):
        signal.receivers = self.stashed_signals.get(signal, [])
        del self.stashed_signals[signal]


if __name__ == '__main__':
    # as decorator
    @disable_signals()
    def func():
        print('signals are disabled 1')


    func()

    # as contextmanager
    with disable_signals():
        print('signals are disabled 2')