Minor code tidying

2 files changed, 62 insertions, 59 deletions

diff --git a/pblive/__init__.py b/pblive/__init__.py
index 149a979..149a71d 100644
--- a/pblive/__init__.py
+++ b/pblive/__init__.py
@@ -18,7 +18,11 @@ import eventlet
 import flask
 import flask_socketio
 
-import pblive.data
+# idk why it doesn't work...
+try:
+	from . import data
+except ImportError:
+	import data
 
 import os
 import random
@@ -29,14 +33,14 @@ import yaml
 eventlet.monkey_patch()
 
 app = flask.Flask('pblive')
-app.jinja_env.globals['data'] = pblive.data
+app.jinja_env.globals['data'] = data
 
 socketio = flask_socketio.SocketIO(app)
 
 # Get server IP address
 tmp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 tmp_socket.connect(('8.8.8.8', 0)) # Connecting to a UDP socket sends no packets
-pblive.data.server_ip = tmp_socket.getsockname()[0]
+data.server_ip = tmp_socket.getsockname()[0]
 tmp_socket.close()
 
 # Load session data
@@ -44,7 +48,7 @@ for f in os.listdir('data'):
 	if f.endswith('.yaml') and not f.startswith('.'):
 		session_name = f[:-5]
 		with open(os.path.join('data', f)) as fh:
-			pblive.data.sessions[session_name] = pblive.data.Session.from_dict(yaml.load(fh), session_name)
+			data.sessions[session_name] = data.Session.from_dict(yaml.load(fh), session_name)
 
 @app.route('/')
 def index():
@@ -56,7 +60,7 @@ def admin():
 
 @app.route('/session/<session_name>')
 def session(session_name):
-	return flask.render_template('session.html', session=pblive.data.sessions[session_name])
+	return flask.render_template('session.html', session=data.sessions[session_name])
 
 @app.route('/image/<location>')
 def image(location):
@@ -65,7 +69,7 @@ def image(location):
 
 @app.route('/admin/session/<session_name>')
 def admin_session(session_name):
-	return flask.render_template('admin_session.html', session=pblive.data.sessions[session_name])
+	return flask.render_template('admin_session.html', session=data.sessions[session_name])
 
 @app.route('/debug')
 def debug():
@@ -75,9 +79,9 @@ def debug():
 def socket_join(session_name):
 	app.logger.debug('New client {} connected'.format(flask.request.sid))
 	
-	session = pblive.data.sessions[session_name]
-	user = pblive.data.User(sid=flask.request.sid, session=session)
-	pblive.data.users[flask.request.sid] = user
+	session = data.sessions[session_name]
+	user = data.User(sid=flask.request.sid, session=session)
+	data.users[flask.request.sid] = user
 	
 	# Send initial colour picker
 	flask_socketio.emit('update', flask.render_template('colour_picker.html', session=session), room=flask.request.sid)
@@ -96,11 +100,11 @@ def render_sidebar(user, session):
 def socket_join(session_name):
 	app.logger.debug('New admin {} connected'.format(flask.request.sid))
 	
-	session = pblive.data.sessions[session_name]
-	user = pblive.data.Admin(sid=flask.request.sid, session=session)
-	pblive.data.admins_lock.acquire()
-	pblive.data.admins[flask.request.sid] = user
-	pblive.data.admins_lock.release()
+	session = data.sessions[session_name]
+	user = data.Admin(sid=flask.request.sid, session=session)
+	data.admins_lock.acquire()
+	data.admins[flask.request.sid] = user
+	data.admins_lock.release()
 	
 	# Send initial screen
 	flask_socketio.emit('update', render_question_admin(session, session.question_num), room=flask.request.sid)
@@ -110,42 +114,42 @@ def socket_join(session_name):
 def socket_disconnect():
 	app.logger.debug('Client {} disconnected'.format(flask.request.sid))
 	
-	if flask.request.sid in pblive.data.users:
-		user = pblive.data.users[flask.request.sid]
+	if flask.request.sid in data.users:
+		user = data.users[flask.request.sid]
 		
 		# Release the colour if it's being held
 		if user.colour:
 			user.session.colours.append(user.colour)
 			
 			# Relay change
-			for _, other_user in pblive.data.iterate_users():
+			for _, other_user in data.iterate_users():
 				if other_user != user and other_user.session == user.session:
 					flask_socketio.emit('update_left', render_sidebar(other_user, user.session), room=other_user.sid)
 					if not other_user.colour:
 						flask_socketio.emit('update', flask.render_template('colour_picker.html', session=user.session), room=other_user.sid)
-			for _, admin in pblive.data.iterate_admins():
+			for _, admin in data.iterate_admins():
 				if admin.session == user.session:
 					flask_socketio.emit('update_left', render_sidebar(admin, user.session), room=admin.sid)
 		
-		pblive.data.users_lock.acquire()
-		del pblive.data.users[flask.request.sid]
-		pblive.data.users_lock.release()
+		data.users_lock.acquire()
+		del data.users[flask.request.sid]
+		data.users_lock.release()
 
 @socketio.on('register')
 def socket_register(colour_id, colour_name):
-	user = pblive.data.users[flask.request.sid]
+	user = data.users[flask.request.sid]
 	
 	if not user.colour and (colour_id, colour_name) in user.session.colours:
 		user.colour = (colour_id, colour_name)
 		user.session.colours.remove(user.colour)
 		
 		# Relay change
-		for _, other_user in pblive.data.iterate_users():
+		for _, other_user in data.iterate_users():
 			if other_user != user and other_user.session == user.session:
 				flask_socketio.emit('update_left', render_sidebar(other_user, user.session), room=other_user.sid)
 				if not other_user.colour:
 					flask_socketio.emit('update', flask.render_template('colour_picker.html', session=user.session), room=other_user.sid)
-		for _, admin in pblive.data.iterate_admins():
+		for _, admin in data.iterate_admins():
 			if admin.session == user.session:
 				flask_socketio.emit('update_left', render_sidebar(admin, user.session), room=admin.sid)
 	
@@ -154,39 +158,39 @@ def socket_register(colour_id, colour_name):
 
 @socketio.on('answer')
 def socket_answer(question_num, answer):
-	user = pblive.data.users[flask.request.sid]
+	user = data.users[flask.request.sid]
 	
 	if question_num == user.session.question_num:
-		if isinstance(user.session.questions[user.session.question_num], pblive.data.SpeedQuestion):
+		if isinstance(user.session.questions[user.session.question_num], data.SpeedQuestion):
 			if question_num in user.answers:
 				# Only one shot!
 				return
 		
 		user.answers[question_num] = answer
 		
-		if isinstance(user.session.questions[user.session.question_num], pblive.data.MCQQuestion):
+		if isinstance(user.session.questions[user.session.question_num], data.MCQQuestion):
 			flask_socketio.emit('update', render_question(user, user.session, user.session.question_num), room=user.sid)
 		
 		# Hurry!
-		if isinstance(user.session.questions[user.session.question_num], pblive.data.SpeedQuestion):
+		if isinstance(user.session.questions[user.session.question_num], data.SpeedQuestion):
 			if user.session.questions[user.session.question_num].timer_thread is None:
-				user.session.questions[user.session.question_num].timer_thread = pblive.data.SpeedQuestionTimerThread(flask.copy_current_request_context(do_goto_question), user.session, user.session.question_num + 1)
+				user.session.questions[user.session.question_num].timer_thread = data.SpeedQuestionTimerThread(flask.copy_current_request_context(do_goto_question), user.session, user.session.question_num + 1)
 				user.session.questions[user.session.question_num].timer_thread.start()
 		
 		# Relay change
-		for _, other_user in pblive.data.iterate_users():
+		for _, other_user in data.iterate_users():
 			if other_user.session == user.session:
 				flask_socketio.emit('update_left', render_sidebar(other_user, user.session), room=other_user.sid)
-				if isinstance(user.session.questions[user.session.question_num], pblive.data.SpeedQuestion):
+				if isinstance(user.session.questions[user.session.question_num], data.SpeedQuestion):
 					flask_socketio.emit('update', render_question(other_user, user.session, user.session.question_num), room=other_user.sid)
-		for _, admin in pblive.data.iterate_admins():
+		for _, admin in data.iterate_admins():
 			if admin.session == user.session:
 				flask_socketio.emit('update', render_question_admin(user.session, user.session.question_num), room=admin.sid)
 				flask_socketio.emit('update_left', render_sidebar(admin, user.session), room=admin.sid)
 
 @socketio.on('reveal_answers')
 def socket_reveal_answers(question_num):
-	user = pblive.data.admins[flask.request.sid]
+	user = data.admins[flask.request.sid]
 	
 	user.session.questions[question_num].revealed = True
 	
@@ -194,46 +198,46 @@ def socket_reveal_answers(question_num):
 
 def do_goto_question(session, question_num):
 	# Cleanup old question
-	if isinstance(session.questions[session.question_num], pblive.data.SpeedQuestion):
+	if isinstance(session.questions[session.question_num], data.SpeedQuestion):
 		if session.questions[session.question_num].timer_thread is not None:
 			session.questions[session.question_num].timer_thread.stop()
 	
 	session.question_num = question_num
 	
 	# Do work for some questions
-	if isinstance(session.questions[question_num], pblive.data.RandomQuestion):
-		session.questions[question_num].answerer = random.choice([other_user for _, other_user in pblive.data.users.items() if other_user.session == session and other_user.colour])
+	if isinstance(session.questions[question_num], data.RandomQuestion):
+		session.questions[question_num].answerer = random.choice([other_user for _, other_user in data.users.items() if other_user.session == session and other_user.colour])
 	
 	# Relay change
-	for _, other_user in pblive.data.iterate_users():
+	for _, other_user in data.iterate_users():
 		if other_user.session == session and other_user.colour:
 			flask_socketio.emit('update', render_question(other_user, session, session.question_num), room=other_user.sid)
 			flask_socketio.emit('update_left', render_sidebar(other_user, session), room=other_user.sid)
-	for _, admin in pblive.data.iterate_admins():
+	for _, admin in data.iterate_admins():
 		if admin.session == session:
 			flask_socketio.emit('update', render_question_admin(session, session.question_num), room=admin.sid)
 			flask_socketio.emit('update_left', render_sidebar(admin, session), room=admin.sid)
 
 @socketio.on('goto_question')
 def socket_goto_question(question_num):
-	user = pblive.data.admins[flask.request.sid]
+	user = data.admins[flask.request.sid]
 	
 	do_goto_question(user.session, question_num)
 
 @socketio.on('pass_question')
 def socket_pass_question():
-	user = pblive.data.admins[flask.request.sid] if flask.request.sid in pblive.data.admins else pblive.data.users[flask.request.sid]
+	user = data.admins[flask.request.sid] if flask.request.sid in data.admins else data.users[flask.request.sid]
 	
-	if isinstance(user.session.questions[user.session.question_num], pblive.data.RandomQuestion):
+	if isinstance(user.session.questions[user.session.question_num], data.RandomQuestion):
 		# Re-randomise answerer
-		user.session.questions[user.session.question_num].answerer = random.choice([other_user for _, other_user in pblive.data.users.items() if other_user.session == user.session and other_user.colour])
+		user.session.questions[user.session.question_num].answerer = random.choice([other_user for _, other_user in data.users.items() if other_user.session == user.session and other_user.colour])
 		
 		# Relay change
-		for _, other_user in pblive.data.iterate_users():
+		for _, other_user in data.iterate_users():
 			if other_user.session == user.session and other_user.colour:
 				flask_socketio.emit('update', render_question(other_user, other_user.session, other_user.session.question_num), room=other_user.sid)
 				flask_socketio.emit('update_left', render_sidebar(other_user, user.session), room=other_user.sid)
-		for _, admin in pblive.data.iterate_admins():
+		for _, admin in data.iterate_admins():
 			if admin.session == user.session:
 				flask_socketio.emit('update', render_question_admin(admin.session, admin.session.question_num), room=admin.sid)
 				flask_socketio.emit('update_left', render_sidebar(admin, user.session), room=admin.sid)
diff --git a/pblive/data.py b/pblive/data.py
index 8d1f2d8..9316331 100644
--- a/pblive/data.py
+++ b/pblive/data.py
@@ -41,13 +41,11 @@ class Session:
 		return cls(name=name, title=obj['title'], questions=[Question.from_dict(x) for x in obj['questions']])
 
 class Question:
-	def __init__(self, prompt=None, image=None, answers=None, revealed=False):
-		if answers is None:
-			answers = []
+	def __init__(self, *args, **kwargs):
+		self.prompt = kwargs.get('prompt', None)
+		self.image = kwargs.get('image', None)
+		self.answers = kwargs.get('answers', [])
 		
-		self.prompt = prompt
-		self.image = image
-		self.answers = answers
 		self.revealed = False
 	
 	@staticmethod
@@ -68,12 +66,9 @@ class Question:
 	def load_dict(self, obj):
 		self.type = obj['type']
 		
-		if 'prompt' in obj:
-			self.prompt = obj['prompt']
-		if 'image' in obj:
-			self.image = obj['image']
-		if 'answers' in obj:
-			self.answers = obj['answers']
+		self.prompt = obj.get('prompt', self.prompt)
+		self.image = obj.get('image', self.image)
+		self.answers = obj.get('answers', self.answers)
 
 class LandingQuestion(Question):
 	# Not actually a question
@@ -92,8 +87,10 @@ class RandomQuestion(Question):
 	template = 'question_random.html'
 	template_admin = 'question_random_admin.html'
 	
-	def __init__(self):
-		self.answerer = None
+	def __init__(self, *args, **kwargs):
+		super().__init__(*args, **kwargs)
+		
+		self.answerer = kwargs.get('answerer', None)
 
 class TypeQuestion(Question):
 	template = 'question_type.html'
@@ -113,7 +110,9 @@ class SpeedQuestion(MCQQuestion):
 	template = 'question_speed.html'
 	template_admin = 'question_speed_admin.html'
 	
-	def __init__(self):
+	def __init__(self, *args, **kwargs):
+		super().__init__(*args, **kwargs)
+		
 		self.timer_thread = None
 
 class SpeedQuestionTimerThread(threading.Thread):