175e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chenimport signal
275e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chenimport weakref
375e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
475e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chenfrom unittest2.compatibility import wraps
575e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
675e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen__unittest = True
775e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
875e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
975e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chenclass _InterruptHandler(object):
1075e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    def __init__(self, default_handler):
1175e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        self.called = False
1275e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        self.default_handler = default_handler
1375e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
1475e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    def __call__(self, signum, frame):
1575e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        installed_handler = signal.getsignal(signal.SIGINT)
1675e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        if installed_handler is not self:
1775e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            # if we aren't the installed handler, then delegate immediately
1875e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            # to the default handler
1975e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            self.default_handler(signum, frame)
2075e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
2175e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        if self.called:
2275e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            self.default_handler(signum, frame)
2375e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        self.called = True
2475e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        for result in _results.keys():
2575e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            result.stop()
2675e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
2775e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen_results = weakref.WeakKeyDictionary()
2875e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chendef registerResult(result):
2975e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    _results[result] = 1
3075e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
3175e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chendef removeResult(result):
3275e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    return bool(_results.pop(result, None))
3375e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
3475e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen_interrupt_handler = None
3575e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chendef installHandler():
3675e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    global _interrupt_handler
3775e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    if _interrupt_handler is None:
3875e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        default_handler = signal.getsignal(signal.SIGINT)
3975e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        _interrupt_handler = _InterruptHandler(default_handler)
4075e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        signal.signal(signal.SIGINT, _interrupt_handler)
4175e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
4275e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
4375e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chendef removeHandler(method=None):
4475e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    if method is not None:
4575e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        @wraps(method)
4675e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        def inner(*args, **kwargs):
4775e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            initial = signal.getsignal(signal.SIGINT)
4875e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            removeHandler()
4975e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            try:
5075e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen                return method(*args, **kwargs)
5175e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen            finally:
5275e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen                signal.signal(signal.SIGINT, initial)
5375e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        return inner
5475e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen
5575e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    global _interrupt_handler
5675e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen    if _interrupt_handler is not None:
5775e28f942c1b9f9c6d5a0d5f2efd037cbbc9fc74Johnny Chen        signal.signal(signal.SIGINT, _interrupt_handler.default_handler)
58