from Queue import Queue
import threading
import traceback
import inspect
import os, stat

class stack_size:
	def __init__(self, stack_size):
		self.stack_size = stack_size
		self.old_size = threading.stack_size()

	def __enter__(self):
		threading.stack_size(self.old_size)

	def __exit__(self, type, value, traceback):
		threading.stack_size(self.stack_size)

class Lock:
	def __init__(self):
		self.lock = threading.Lock()
	
	def __enter__(self):
		self.lock.acquire()
	
	def __exit__(self, type, value, traceback):
		self.lock.release()

class FilenameLogger:
	def __init__(self, filename):
		self.filename = filename
		self.filepath = os.path.join(os.getcwd(), filename)
		self.fileobj = open(self.filepath, 'a')
	
	def __call__(self, msg):
		# someone renamed our file! maybe logrotate'd. open a new one.
		if not os.path.exists(self.filepath) or os.stat(self.filepath)[stat.ST_INO] != os.fstat(self.fileobj.fileno())[stat.ST_INO]:
			self.fileobj = open(self.filepath, 'a')
		
		self.fileobj.write(msg+'\n')

class FileObjLogger:
	def __init__(self, fileobj=None, filename=None):
		self.fileobj = fileobj or open(filename, 'a') if filename else None
	
	def __call__(self, msg):
		if self.fileobj and not self.fileobj.closed:
			self.fileobj.write(msg+'\n')

class ThreadSafeLogger:
	'''
	Init with a list of loggers to call when this class is called()
	"None" logs to stdout, string obj logs to filename, fileobj logs to write(), callable(obj) acts as a callback.

	formatter=callback allows you to set a callback which can rewrite the message passed to log
	separator=str allows you to set the separator between args passed
	'''
	def __init__(self, *args, **kwargs):
		self.formatter = None
		self.separator = ', '
		if 'formatter' in kwargs and callable(kwargs['formatter']):
			self.formatter = kwargs['formatter']
		if 'separator' in kwargs:
			self.separator = kwargs['separator']
		
		loggers = set()
		for logger in args:
			if hasattr(logger, 'write') and callable(logger.write):
				loggers.add(FileObjLogger(fileobj=logger))
			elif isinstance(logger, basestring):
				loggers.add(FilenameLogger(logger))
			elif callable(logger):
				loggers.add(logger)
			else:
				loggers.add(self.echo)
		
		self.loggers = loggers or (self.echo,)
		self.lock = Lock()
	
	def __call__(self, *args):
		with self.lock:
			self.callback(*args)

	def callback(self, *args):
		msg = self.separator.join(str(arg) for arg in args)
		if self.formatter:
			msg = self.formatter(msg)

		for logger in self.loggers:
			logger(msg)

	def echo(self, msg):
		print msg

class ThreadPool:
	def __init__(self, max_threads, log_returns=False, catch_returns=False, logger=None, stack_size=0):
		self.lock = threading.Lock()
		self.max = max_threads
		self.logger = logger
		self.stack_size = stack_size
		self.log_returns = log_returns
		self.catch_returns = catch_returns

		self.call_queue = Queue()
		self.returns = Queue()
		self.spawn_workers()

	def __call__(self, func):
		def wrapper(*args, **kwargs):
			self.call_queue.put((func, args, kwargs))

		return wrapper
	
	def spawn_workers(self):
		for i in xrange(self.max):
			thread = threading.Thread(target=self.worker, args=(self.call_queue,))
			thread.daemon = True
			thread.start()
	
	def worker(self, call):
		while True:
			func, args, kwargs = call.get()
			try:
				result = func(*args, **kwargs)
				if self.catch_returns or self.log_returns:
					if result:
						if isinstance(result, dict):
							self.returned(result)
						elif inspect.isgenerator(result):
							for x in result:
								self.returned(x)
						else:
							self.returned(result)
			except:
				self.logger(traceback.format_exc())
			finally:
				call.task_done()
	
	def returned(self, result):
		if self.log_returns:
			self.logger(result)
		if self.catch_returns:
			self.returns.put(result)
	
	def pop(self):
		'''
		pop a result from the queue, blocks if we have none
		'''
		if self.catch_returns:
			result = self.returns.get()
			self.returns.task_done()
			return result

	def flush(self):
		'''
		clear and return the function returns queue
		'''
		if self.catch_returns:
			results = tuple(self.returns.queue)
			self.returns = Queue()
			return results

		return ()

	def finish(self):
		'''
		wait for queue to finish, then return flush()
		'''
		self.call_queue.join()
		return self.flush()

if __name__ == '__main__':
	log = ThreadSafeLogger()
	# don't use catch_returns unless you're going to use them - otherwise they will never clear and just leak memory
	pool = ThreadPool(100, logger=log, log_returns=True, catch_returns=True)

	import time

	@pool
	def test(i):
		log('from thread', i)
		yield 'yield %i' % i
		# sleep to show how awesome it is that we allow generators
		time.sleep(0.5)
		yield 'yield %i' % (i*10)

	for i in xrange(1,6):
		test(i)

	# because these pops will all happen before the sleep is finished, we'll get numbers < 100
	log('first pop')
	for i in xrange(5):
		log(pool.pop())
	
	# just to make sure it works
	pool.flush()

	# because these pops will happen after the sleep, we'll get numbers over 100
	results = pool.finish()
	print results
@@ -2749,16 +2749,37 @@
 = logger
+ or (lambda *x: None)
 %0A%09%09self.
@@ -5117,16 +5117,17 @@
 )%0A%09print results
+%0A

fix crash on exceptions with no logger