Decorator / Context manager to wait and notify on conditions

Condition Decorators

from functools import wraps
from threading import Condition

"""
Usage: 

cond = Condition()

Acquire the condition, wait for it, then do something: 
@WaitOn(cond)
def f():
    something()

or 

def f():
    with WaitOn();
        something()

Acquire the condition, do something, then notify: 
@NotifyOn(cond)
def g():
    something()

or 

def g():
    with NotifyOn(cond):
        something()
"""


class WaitOn:
    def __init__(self, cond):
        self.cond = cond
    def __enter__(self):
        self.cond.acquire()
        self.cond.wait()
        return self
    def __exit__(self, type, value, traceback):
        self.cond.release()
    def __call__(self, fn):
        @wraps(fn)
        def deco(*args, **kwargs):
            with self:
                return fn(*args, **kwargs)
        return deco

class NotifyOn:
    def __init__(self, cond):
        self.cond = cond
    def __enter__(self):
        self.cond.acquire()
        return self
    def __exit__(self, type, value, traceback):
        self.cond.notify_all()
        self.cond.release()
    def __call__(self, fn):
        @wraps(fn)
        def deco(*args, **kwargs):
            with self:
                return fn(*args, **kwargs)
        return deco