Django uwsgi runner with reloader

run-django.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import unicode_literals

import argparse
import atexit
import collections
import itertools
import os
import signal
import socket
import subprocess
import sys
import threading

from watchdog import events
from watchdog.observers import Observer


try:
	import configparser
except ImportError:
	import ConfigParser as configparser


LOG_SEPARATOR = b'\t'
LOG_VARS = ['uri', 'method', 'addr', 'status', 'micros', 'msecs', 'time', 'size', ]
METHOD_COLORS = {b'GET': 32, b'POST': 33, b'DEFAULT': 36}
STATUS_COLORS = {b'200': (32, 1), b'301': (33, 1), b'302': (33, 1), b'404': (31, 1), b'500': (35, 1), b'DEFAULT': (37, 0)}


LogLine = collections.namedtuple('LogLine', LOG_VARS)


class ReloaderEventHandler(events.PatternMatchingEventHandler):
	RELOAD_ON_EVENTS = {
		events.EVENT_TYPE_MOVED,
		events.EVENT_TYPE_DELETED,
		events.EVENT_TYPE_CREATED,
		events.EVENT_TYPE_MODIFIED,
	}

	def __init__(self, *args, **kwargs):
		self.proc = kwargs.pop('proc')
		self.reload_wait_time = kwargs.pop('reload_wait_time')
		self.timer = None
		super(ReloaderEventHandler, self).__init__(*args, **kwargs)

	def _run(self):
		if self.timer is not None:
			self.timer.cancel()
		pid = self.proc.pid
		def reload_uwsgi():
			print("Reloading uwsgi ...")
			os.kill(pid, signal.SIGHUP)
			self.timer = None
		self.timer = threading.Timer(self.reload_wait_time, reload_uwsgi)
		self.timer.start()

	def on_any_event(self, event):
		if event.event_type in self.RELOAD_ON_EVENTS:
			self._run()


class ExtraCommandHandler(events.PatternMatchingEventHandler):
	RELOAD_ON_EVENTS = {
		events.EVENT_TYPE_MOVED,
		events.EVENT_TYPE_DELETED,
		events.EVENT_TYPE_CREATED,
		events.EVENT_TYPE_MODIFIED,
	}

	def __init__(self, *args, **kwargs):
		self.busy = False
		self.command = kwargs.pop('command')
		super(ExtraCommandHandler, self).__init__(*args, **kwargs)

	def _run(self):
		if self.busy:
			return
		self.busy = True
		os.system(self.command)
		def unlock():
			self.busy = False
		timer = threading.Timer(1, unlock)
		timer.start()

	def on_any_event(self, event):
		if event.event_type in self.RELOAD_ON_EVENTS:
			self._run()


class ConfigError(RuntimeError):
	pass


def flush_memcache():
	sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	try:
		sock.connect(('localhost', 11211))
	except Exception:
		return
	sock.send(b'flush_all\nquit\n')
	sock.close()


def main():
	flush_memcache()
	settings = get_settings()
	#default_params = '--master --cheap --idle 300'
	#default_params = '--master --lazy-apps'
	#default_params = '--master --cheap'
	default_params = '--master'
	params = [settings['uwsgi_binary']] + default_params.split()
	params += ['--plugin', settings['plugin']]
	if settings['listen'].startswith('socket:'):
		params += ['--socket', settings['listen'][len('socket:'):]]
	else:
		params += ['--http', settings['listen']]
	params += ['--module', settings['module']]
	params += ['--virtualenv', settings['virtualenv']]
	params += ['--logformat', settings['logformat']]
	params += ['--buffer-size', '32768']
	params += ['--enable-threads']
	if settings['ini']:
		params += ['--ini', settings['ini']]
	params += settings['forward_args']
	print('Running %s' % ' '.join("'" + p + "'" for p in params))
	try:
		proc = subprocess.Popen(params, stderr=subprocess.PIPE, preexec_fn=os.setpgrp)
		group = os.getpgid(proc.pid)

		def at_exit(*args):
			proc.terminate()
			try:
				proc.wait(0.1)
			except:
				proc.kill()
			try:
				os.killpg(group, signal.SIGKILL)
			except Exception:
				pass
			sys.exit()
		atexit.register(at_exit)
		signal.signal(signal.SIGTERM, at_exit)
		signal.signal(signal.SIGINT, at_exit)
		signal.signal(signal.SIGHUP, at_exit)

		event_handler = ReloaderEventHandler(patterns=['*.py', '*.mo'], proc=proc, reload_wait_time=settings['reload_wait_time'])
		observer = Observer()

		def register_links(handler, directory):
			for filename in os.listdir(directory):
				path = os.path.join(directory, filename)
				if os.path.islink(path):
					path = os.path.abspath(os.readlink(path))
					if os.path.isdir(path):
						observer.schedule(handler, path=path, recursive=True)
				if os.path.isdir(path):
					register_links(handler, path)

		for pattern, command in settings['extra_commands'].items():
			command_handler = ExtraCommandHandler(patterns=pattern.split(','), command=command)
			observer.schedule(command_handler, path=os.path.abspath("."), recursive=True)
			register_links(command_handler, ".")

		observer.schedule(event_handler, path=os.path.abspath("."), recursive=True)
		register_links(event_handler, ".")
		observer.start()

		output = getattr(sys.stdout, 'buffer', sys.stdout)
		for line in iter(proc.stderr.readline, b''):
			try:
				log_line = LogLine(*line.split(LOG_SEPARATOR))
			except TypeError:
				output.write(line)
				output.flush()
				continue
			parts = [
				colorize(pad(log_line.method, 5), METHOD_COLORS.get(log_line.method, METHOD_COLORS[b'DEFAULT']), 1),
				colorize(pad_r(log_line.msecs, 4) + b' ms', 34, 1),
				colorize(log_line.status, *STATUS_COLORS.get(log_line.status, STATUS_COLORS[b'DEFAULT'])),
				log_line.uri,
			]
			output.write(b' '.join(parts) + b'\n')
			output.flush()
	except KeyboardInterrupt:
		output.write(proc.stderr.read())
		output.flush()



def read_settings_file():
	global_path = os.path.join(os.environ.get('XDG_CONFIG_HOME', '~/.config'), 'run_django.cfg')
	config = configparser.ConfigParser()
	config.read([global_path, 'run_django.cfg'])
	settings = {
		'plugin': 'python%d%d' % (sys.version_info.major, sys.version_info.minor),
		'listen': '127.0.0.1:8000',
		'module': None,
		'uwsgi_binary': 'uwsgi',
		'virtualenv': os.environ.get('VIRTUAL_ENV'),
		'logformat': LOG_SEPARATOR.decode('utf-8').join('%(' + logvar + ')' for logvar in LOG_VARS),
		'reload_wait_time': 0.1,
		'ini': None,
	}
	for key, value in config.items('DEFAULT'):
		#if not key in settings:
		#	raise ConfigError("Settings key '%s' not defined" % key)
		settings[key] = value
	settings['extra_commands'] = {}
	for section in config.sections():
		settings['extra_commands'][section] = config.get(section, 'command')
	return settings


def get_settings():
	settings = read_settings_file()
	parser = argparse.ArgumentParser()
	parser.add_argument(
		'--plugin',
		type=str,
		help="uwsgi plugin, default %s" % settings['plugin']
	)
	parser.add_argument(
		'--listen',
		type=str,
		help="address, default %s" % settings['listen']
	)
	parser.add_argument(
		'--module',
		type=str,
		help="python module to run"
	)
	parser.add_argument(
		'--uwsgi_binary',
		type=str,
		help="path to uwsgi binary, default %s" % settings['uwsgi_binary']
	)
	parser.add_argument(
		'--virtualenv',
		type=str,
		help="path to uwsgi virtual environment, default %s" % settings['virtualenv'] if settings['virtualenv'] else "path to uwsgi virtual environment"
	)
	parser.add_argument(
		'--logformat',
		type=str,
		help="uwsgi log format"
	)
	parser.add_argument(
		'--reload_wait_time',
		type=float,
		help="wait time before reload"
	)
	parser.add_argument(
		'--ini',
		type=str,
		help="path to ini file"
	)
	args = parser.parse_args(list(itertools.takewhile(lambda arg: arg != '--', sys.argv[1:])))
	settings['forward_args'] = list(itertools.dropwhile(lambda arg: arg != '--', sys.argv[1:]))[1:]
	settings['plugin'] = args.plugin or settings['plugin']
	settings['listen'] = args.listen or settings['listen']
	settings['module'] = args.module or settings['module']
	settings['uwsgi_binary'] = args.uwsgi_binary or settings['uwsgi_binary']
	settings['logformat'] = args.logformat or settings['logformat']
	settings['reload_wait_time'] = args.reload_wait_time or settings['reload_wait_time']
	settings['ini'] = args.ini or settings['ini']
	if not settings['module']:
		raise ConfigError("Argument '%s' is required" % 'module')
	if not settings['virtualenv']:
		raise ConfigError("Argument '%s' is required" % 'virtualenv')
	return settings


def colorize(text, color, light=0):
	start = b'\033[' + str(light).encode('utf-8') + b';' + str(color).encode('utf-8') + b'm'
	end = b'\033[0;0m'
	return start + text + end


def pad(text, width):
	if len(text) >= width:
		return text
	else:
		return text + b' ' * (width - len(text))


def pad_r(text, width):
	if len(text) >= width:
		return text
	else:
		return b' ' * (width - len(text)) + text


if __name__ == "__main__":
	try:
		main()
	except ConfigError as e:
		print(e)