
from __future__ import absolute_import
import threading


class RWLock(object):

    def __init__(self):
        self.shared = Context(self.acquire_read, self.release)
        self.exclusive = Context(self.acquire_write, self.release)
        self._lock = threading.Lock()
        self._waiters = []
        self._holders = {}
        self._writer = None

    def acquire_write(self):
        me = threading.current_thread()
        if me is self._writer:
            self._holders[me] += 1
            return
        if me in self._holders:
            raise RuntimeError("Lock promotion is forbidden")
        with self._lock:
            if self._holders or self._waiters:
                self._wait(True)
            self._holders[me] = 1
            self._writer = me

    def acquire_read(self):
        me = threading.current_thread()
        if me is self._writer:
            raise RuntimeError("Lock demotion is forbidden")
        if me in self._holders:
            self._holders[me] += 1
            return
        with self._lock:
            if self._writer or self._waiters:
                self._wait(False)
            self._holders[me] = 1
            if self._waiters:
                self._grant_next_waiter()

    def release(self):
        me = threading.current_thread()
        if me not in self._holders:
            raise RuntimeError("Thread %s attempted to release a lock it "
                               "does not hold" % me)
        self._holders[me] -= 1
        if self._holders[me] > 0:
            return
        with self._lock:
            self._writer = None
            del self._holders[me]
            if self._waiters:
                self._grant_next_waiter()

    def _wait(self, wants_write):
        waiter = Waiter(wants_write)
        self._waiters.append(waiter)
        try:
            self._lock.release()
            try:
                waiter.wait()
            finally:
                self._lock.acquire()
        finally:
            self._waiters.remove(waiter)

    def _grant_next_waiter(self):
        if self._holders and self._waiters[0].wants_write:
            return
        self._waiters[0].grant()


class Waiter(object):

    def __init__(self, wants_write):
        self.wants_write = wants_write
        self._event = threading.Event()

    def wait(self):
        self._event.wait()

    def grant(self):
        self._event.set()


class Context(object):

    def __init__(self, acquire, release):
        self._acquire = acquire
        self._release = release

    def __enter__(self):
        self._acquire()
        return self

    def __exit__(self, *args):
        self._release()
