from __future__ import with_statement

"""This module contains several variants on mailboxes.
A mailbox is a way to pass information from one thread to another.
You put something into a mailbox; eventually the other thread takes
it out.

These mailboxes  also have mechanisms so that you can shut down
a mailbox.   You can "close" it so that the other side will
get an exception when it tries to read or write.

A typical server looks like this::

	mbox = mailbox()

	... fork the client thread ...

	while True:
		try:
			x = mbox.get()
		except g_mailgox.EOF:
			break
		serve(x)

A typical client::

	mbox = mailbox()

	... fork the server thread ...

	for task in tasklist:
		mbox.put(task)
	mbox.putclose()	# This causes the next "get" to raise an EOF.
"""

import threading as T
import sys
Lock = T.Lock
Condition = T.Condition

DEBUG = False
# DEBUG = True

if __debug__:
	_prinfo_lock = T.Lock()
	def prinfo(*s):
		global _prinfo_lock
		if not DEBUG:
			return
		with _prinfo_lock:
			sys.stderr.flush()
			sys.stdout.flush()
			tid = '???'
			sys.stderr.writelines([tid] + [ str(q) for q in s] + ['\n'])
			sys.stderr.flush()
else:
	def prinfo(*s):
		pass



def _Raise(x):
	if isinstance(x, (Exception, SystemExit, KeyboardInterrupt)):
		# print 'Exception', x
		raise x
	elif isinstance(x, tuple):
		# print 'tuple', x
		if len(x) == 1:
			raise x[0]
		elif len(x) == 2:
			raise x[0], x[1]
		elif len(x) == 3:
			raise x[0], x[1], x[2]
	raise ValueError, "_Raise needs an exception or a tuple: got %s: %s" % (type(x), str(x))


class EOF(Exception):
	def __init__(self, *s):
		Exception.__init__(self, *s)


class mailbox(object):
	_empty = EOF()
	eof = EOF()
	ignore = EOF()

	def __init__(self):
		self.lock = T.Lock()
		self.gc = T.Condition(self.lock)
		self.pc = T.Condition(self.lock)
		self.v = self._empty
		self.psui = None
		self.gsui = None
		self.nps = -1
		self.gps = -1


	def close(self, for_put=eof, for_get=eof, ngx=-1, npx=-1):
		"""Close the connection.
		It causes other threads that use the mailbox to raise an exception,
		unless you give it self._ignore as an argument.
		Npx and ngx specify how many threads will commit suicide;
		zero means none; setting either to 1 is a useful way to kill
		a single thread.

		Legal values for for_put and for_get are any exception,
		anything returned from sys.exc_info, self.ignore and None.
			- None => Do not raise an exception;
			- self.ignore => Do not change whatever exception (or None)
				is already stored.
		"""
		with self.lock:
			self.__set_psui(for_put, npx)
			self.__set_gsui(for_get, ngx)


	def __gsuicide(self):
		"""@requires: Called under lock.
		@note: suicide never happens when the mailbox is full.
		"""
		if self.v is self._empty and self.gsui is not None:
			tmpe = self.gsui
			if self.ngs > 0:
				self.ngs -= 1
			if self.ngs == 0:
				self.gsui = None
			self.gc.notify()	# Necessary for getclose(); get() case.
			_Raise(tmpe)

	def __psuicide(self):
		"""@requires: Called under lock."""
		if self.psui is not None:
			tmpe = self.psui
			if self.nps > 0:
				self.nps -= 1
			if self.nps == 0:
				self.psui = None
			self.pc.notify()	# Necessary for putclose(); put() case.
			_Raise(tmpe)


	def __set_psui(self, for_put, npx):
		"""@requires: Called under lock."""
		if for_put is not self.ignore and self.psui is None:
			if npx == 0:
				self.psui = None
			else:
				self.psui = for_put
				self.pc.notifyAll()
			self.nps = npx


	def __set_gsui(self, for_get, ngx):
		"""@requires: Called under lock."""
		if for_get is not self.ignore and self.gsui is None:
			if ngx == 0:
				self.gsui = None
			else:
				self.gsui = for_get
				self.gc.notifyAll()
			self.ngs = ngx


	def __len__(self):
		with self.lock:
			tmp = int( self.v is self._empty )
		return tmp

	def __nonzero__(self):
		with self.lock:
			tmp = self.v is self._empty
		return tmp


	def put(self, v):
		"""Put a value into the mailbox.   If one of the close()
		functions has been used to set an exception on for_put,
		it will raise the exception instead.
		"""
		assert v is not self._empty
		if __debug__:
			prinfo('ready to put %s' % str(v))
		with self.pc:
			while self.v is not self._empty and self.psui is None:
				self.pc.wait()
			self.__psuicide()
			self.v = v
			self.gc.notify()
		if __debug__:
			prinfo('put done')


	def putclose(self, for_get=eof, for_put=ignore, ngx=-1, npx=-1, defer_ex=False):
		"""Close the connection.
		This is normally called by the producer of data
		in place of put().  It can cause other threads to raise an exception.
		
		This call acts like a
		put() in that it will block until the mailbox
		is emptied.  It does not fill the mailbox, though.
		Following this call:
			- get() will not block (it will raise an exception
				unless for_get is self.ignore or npx==0),
			- put() will not block (it will raise an exception if
				for_put is not self.ignore and npx!=0),
			- likewise putclose() and getclose().

		It will raise an exception (unless defer_ex is True)
		if some other thread has closed
		first and set a for_put exception.   However, it will
		set a for_get exception (if requested) before expiring.
		"""
		# print 'putclose'
		with self.pc:
			while self.v is not self._empty and self.psui is None:
				self.pc.wait()
			self.__set_gsui(for_get, ngx)
			if not defer_ex:
				self.__psuicide()
			# self.v remains empty
			self.__set_psui(for_put, npx)
			self.pc.notify()	# A put can follow a putclose()


	def putraise(self):
		"""Raise any residual exceptions.
		This is normally called by the producer of data
		
		This call acts like a
		put() in that it will block until the mailbox
		is emptied or until an exception is waiting.
		"""
		# print 'putclose'
		with self.pc:
			while self.v is not self._empty and self.psui is None:
				self.pc.wait()
			self.__psuicide()
			self.pc.notify()	# A put or putclose can follow a putraise()


	def get(self):
		"""Get an item from the mailbox.  Possibly, raise an exception if L{close} or L{putclose} has been called on the mailbox.
		@note: Suicide has a lower priority than processing data.
			So, if there is data in the mailbox and putclose() has
			been called, one call to get() will succeed then
			the second call will raise an exception.
		@note: If the mailbox is empty, and if a previous call has specified
			an exception to be raised by readers, raise it.
		"""
		if __debug__:
			prinfo('ready to get acquire')
		with self.gc:
			while self.v is self._empty and self.gsui is None:
				self.gc.wait()
			tmp = self.v
			self.__gsuicide()
			self.v = self._empty
			self.gc.notify()	# Even though we empty the mailbox, we
						# need to unblock the next call to get()
						# because it might need to suicide 
						# (if a suitable exception is set).
			self.pc.notify()
		if __debug__:
			prinfo('get done=%s' % str(tmp))
		return tmp


	def getclose(self, for_put=eof, for_get=ignore, npx=-1, ngx=-1, defer_ex=False):
		"""Close the connection.
		This is normally called by the consumer of data
		in place of get().
		It causes other threads to raise an exception.
		It will raise an exception if some other thread has closed
		the mailbox first.
		Note that if there is something in the mailbox, it will
		empty and return the value.
		@note: If the mailbox is empty, and if a previous call has specified
			an exception to be raised by readers, raise it.
		"""
		with self.gc:
			while self.v is self._empty and self.gsui is None:
				self.gc.wait()
			tmp = self.v
			self.__set_psui(for_put, npx)
			if not defer_ex:
				self.__gsuicide()
			self.v = self._empty
			self.pc.notify()
			self.__set_gsui(for_get, ngx)
		return tmp



	if __debug__ and DEBUG:
		def __del__(self):
			# We don't need to acquire or release any locks, because we only
			# get to __del__ if there is nothing holding a reference
			# to this object, so that nothing can be waiting on a lock.
			if self.v is not self._empty and not self.gsui and not self.psui:
				prinfo('Mailbox deleted while non-empty')


class waitbox(object):
	"""This is a stripped down mailbox that is good for one message
	between one sender and one reciever.
	It is an error to call put() or get() or putclose() more than once.

	The normal usage is that you create the mailbox, spawn a thread, and hand the mailbox to the new thread.
	Then, one side calls L{get} on the mailbox to wait for an answer.
	The other side computes something, then calls L{put}.
	The waitbox is then thrown away.
	"""
	_empty = EOF()
	_was_closed = EOF()
	_has_been_used = EOF()

	def __init__(self):
		self.get_lock = T.Lock()
		self.get_lock.acquire()
		self.v = self._empty
		self.sui = None

	def __nonzero__(self):
		with self.get_lock:
			tmp = self.v is self._empty
		return tmp

	def putclose(self, to_be_raised=_was_closed):
		"""Close the connection.
		This is normally called by the producer of data in place of L{put}().
		It causes other threads to raise an exception.
		"""
		self.sui = to_be_raised
		self.get_lock.release()


	def put(self, v):
		self.v = v
		self.get_lock.release()


	def get(self):
		"""Get an object from the mailbox.   This will normally block if the mailbox is empty.
		@note: Suicide has a lower priority than processing data.
			So, if there is data in the mailbox and L{putclose}() has
			been called, one call to C{get()} will succeed;
			the second call will raise an exception.
		@note: If the mailbox is empty, and if a previous call has specified
			an exception to be raised by readers, raise it.
		"""
		self.get_lock.acquire()
		self.get_lock = None
		if self.sui is not None:
			_Raise(self.sui)
		self.sui = self._has_been_used
		return self.v


	if __debug__ and DEBUG:
		def __del__(self):
			if self.v is not self._empty and not self.sui:
				prinfo('Mailbox deleted while non-empty')
			# We don't need to unlock the lock, because we only
			# get to __del__ if there is nothing holding a reference
			# to this object, so that nothing can be waiting on a lock.


class repbox(object):
	"""Like a mailbox, but produces infinite copies of whatever value
	was put in there last.   It is created with a value inside.
	@note: This needs revision to more closely follow mailbox's methods.
	"""

	_empty = EOF()

	def __init__(self, value=None, suicide=None):
		self.value = value
		self.suicide = suicide
		self.lock = T.Lock()


	def close(self, e_to_raise=_empty):
		with self.lock:
			if self.suicide is None:
				self.suicide = e_to_raise


	def __len__(self):
		with self.lock:
			if self.suicide is not None:
				tmp = self.suicide
				self.lock.release()
				_Raise(tmp)
		return 1


	def closenow(self, e_to_raise=_empty):
		with self.lock:
			self.suicide = e_to_raise


	def put(self, v):
		with self.lock:
			if self.suicide is not None:
				tmp = self.suicide
				self.lock.release()
				_Raise(tmp)
			self.value = v


	def get(self):
		"""Get an object from the mailbox.   This will normally block if the mailbox is empty.
		@note: If the mailbox is empty, and if a previous call has specified
			an exception to be raised by readers, raise it.
		"""
		with self.lock:
			if self.suicide is not None:
				tmp = self.suicide
				self.lock.release()
				_Raise(tmp)
			tmp = self.value
		return tmp


class listbox(object):
	"""A mailbox with an unbounded queue."""
	eof = EOF()
	ignore = EOF()

	def __init__(self, initial=None):
		self.lock = T.Lock()
		self.gc = T.Condition(self.lock)
		self.pc = self.lock
		if initial:
			self.v = list(initial)
		else:
			self.v = []
		self.psui = None
		self.gsui = None
		self.nps = -1
		self.gps = -1


	def close(self, for_put=eof, for_get=eof, ngx=-1, npx=-1):
		"""Close the connection.
		It causes other threads to raise an exception,
		unless you give it self._ignore as an argument.
		Npx and ngx specify how many threads will commit suicide;
		zero means none; setting either to 1 is a useful way to kill
		a single thread.

		Legal values for for_put and for_get are any exception,
		anything returned from sys.exc_info, self.ignore and None.
			- None => Do not raise an exception;
			- self.ignore => Do not change whatever exception (or None)
				is already stored.
		"""
		with self.lock:
			self.__set_psui(for_put, npx)
			self.__set_gsui(for_get, ngx)


	def __gsuicide(self):
		"""@requires: Called under lock.
		@note: Suicide never happens when the mailbox is full.
		"""
		if (not self.v) and (self.gsui is not None):
			tmpe = self.gsui
			if self.ngs > 0:
				self.ngs -= 1
			if self.ngs == 0:
				self.gsui = None
			self.gc.notify()	# Necessary for getclose(); get() case.
			_Raise(tmpe)

	def __psuicide(self):
		"""@requires: Called under lock."""
		if self.psui is not None:
			tmpe = self.psui
			if self.nps > 0:
				self.nps -= 1
			if self.nps == 0:
				self.psui = None
			_Raise(tmpe)


	def __set_psui(self, for_put, npx):
		"""@requires: Called under lock."""
		if for_put is not self.ignore and self.psui is None:
			if npx == 0:
				self.psui = None
			else:
				self.psui = for_put
			self.nps = npx


	def __set_gsui(self, for_get, ngx):
		"""@requires: Called under lock."""
		if for_get is not self.ignore and self.gsui is None:
			if ngx == 0:
				self.gsui = None
			else:
				self.gsui = for_get
				self.gc.notifyAll()
			self.ngs = ngx


	def __len__(self):
		with self.lock:
			tmp = len(self.v)
		return tmp


	def put(self, v):
		"""Put a value into the mailbox.   If one of the close()
		functions has been used to set an exception on for_put,
		it will raise the exception instead.
		"""
		if __debug__:
			prinfo('ready to put %s' % str(v))
		with self.pc:
			self.__psuicide()
			self.v.append(v)
			self.gc.notify()
		if __debug__:
			prinfo('put done')


	def putclose(self, for_get=eof, for_put=ignore, ngx=-1, npx=-1, defer_ex=False):
		"""Close the connection.
		This is normally called by the producer of data
		in place of put().  It can cause other threads to raise an exception.
		
		It does not fill the mailbox.
		Following this call,
			- get() will not block (it will raise an exception
				unless for_get is self.ignore or npx==0),
			- put() will not block (it will raise an exception if
				for_put is not self.ignore and npx!=0),
			- likewise putclose() and getclose().

		It will raise an exception (unless defer_ex is True)
		if some other thread has closed
		first and set a for_put exception.   However, it will
		set a for_get exception (if requested) before expiring.
		"""
		# print 'putclose'
		with self.pc:
			self.__set_gsui(for_get, ngx)
			if not defer_ex:
				self.__psuicide()
			# self.v remains empty
			self.__set_psui(for_put, npx)


	def putraise(self):
		"""Raise any residual exceptions.
		This is normally called by the producer of data
		"""
		# print 'putclose'
		with self.pc:
			self.__psuicide()


	def get(self):
		"""Get an object from the mailbox.   This will normally block if the mailbox is empty.
		@note: Suicide has a lower priority than processing data.
			So, if there is data in the mailbox and putclose() has
			been called, one call to get() will succeed;
			the second call will raise an exception.
		@note: If the mailbox is empty, and if a previous call has specified
			an exception to be raised by readers, raise it.
		"""
		if __debug__:
			prinfo('ready to get acquire')
		with self.gc:
			while (not self.v) and self.gsui is None:
				self.gc.wait()
			if self.v:
				tmp = self.v.pop(0)
			else:
				self.__gsuicide()
			self.gc.notify()	# Even if we empty the mailbox, we
						# need to unblock the next call to get()
						# because it might need to suicide 
						# (if a suitable exception is set).
		if __debug__:
			prinfo('get done=%s' % str(tmp))
		return tmp


	def getclose(self, for_put=eof, for_get=ignore, npx=-1, ngx=-1, defer_ex=False):
		"""Get an object from the mailbox and close the mailbox.
		This is normally called by the consumer of data
		in place of get().
		It causes other threads to raise an exception.
		@note: It will raise an exception if some other thread has closed the mailbox first.
		"""
		with self.gc:
			while (not self.v) and self.gsui is None:
				self.gc.wait()
			tmp = self.v.pop(0)
			self.__set_psui(for_put, npx)
			if (not self.v) and (not defer_ex):
				self.__gsuicide()
			self.__set_gsui(for_get, ngx)
		return tmp



	if __debug__ and DEBUG:
		def __del__(self):
			# We don't need to acquire or release any locks, because we only
			# get to __del__ if there is nothing holding a reference
			# to this object, so that nothing can be waiting on a lock.
			if self.v and not self.gsui and not self.psui:
				prinfo('Mailbox deleted while non-empty')


maillist = listbox


def test1(m, n=1000):
	for i in range(n):
		m.put(i)
		if m.get() != i:
			print 'Tmp != i'
			sys.exit(1)
	prinfo('OK1')
	m.putclose(KeyError())
	try:
		m.get()
		print 'Whoops1'
		sys.exit(1)
	except KeyError:
		pass	# Got EOF
	prinfo('OK2')

def test2(m):
	m.put(0)
	if m.get() != 0:
		print 'm.get() != 0'
		sys.exit(1)
	m.putclose(KeyError())
	m.put(1)
	# Get will not raise an exception here because
	# there is something in the mailbox.
	if m.get() != 1:
		print 'm.get() != 1'
		sys.exit(1)
	try:
		m.get()
		prinfo('Whoops1')
		sys.exit(1)
	except KeyError:
		pass	# Got EOF

	m.putclose(for_put=EOF())
	prinfo('OK3')
	try:
		m.put(0)
		prinfo('Whoops2')
		sys.exit(3)
	except EOF:
		pass	# Got EOF
	else:
		prinfo('No exception 2')
		sys.exit(4)
	prinfo('OK4')


testmt_waiting = 0
def testmt():

	def innertest(m1, m2, m3, m4):
		global testmt_waiting
		i = m1.get()
		prinfo('innertest: Got m1=%d' % i)
		testmt_waiting += 1
		prinfo('innertest: putting m2=%d' % i)
		m2.put(i)
		i = m3.get()
		testmt_waiting -= 1
		m4.put(i)

	m1 = mailbox()
	m2 = mailbox()
	m3 = mailbox()
	m4 = mailbox()
	prinfo('ok1')
	for i in range(10):
		t = T.Thread(target=innertest, args=(m1,m2,m3,m4))
		t.start()
		prinfo('ok2')
	for i in range(10):
		m1.put(i)
		prinfo('ok3 %d' % i)
		prinfo('ok4, %d' % m2.get())
	if testmt_waiting != 10:
		print 'testmt_waiting != 10'
		sys.exit(1)
	for i in range(10):
		m3.put(i)
		prinfo('ok5 %d' % i)
		prinfo('ok6 %d' % m4.get())
	if testmt_waiting != 0:
		print 'Waiting=%d' % testmt_waiting
		sys.exit(1)

if __name__ == '__main__':
	test1(mailbox(), 30)
	test2(mailbox())
	testmt()
	for i in range(5,100,20):
		for sci in range(0, 100, 10):
			sys.setcheckinterval(sci)
			test1(mailbox(), i)
			test2(mailbox())
			testmt()
	test1(maillist(), 30)
	test2(maillist())
	print "OK"
