Titivillus/quest/consumers.py

#!/usr/bin/env python3
"""
Consumers available for the /quest websocket.
"""
import json

from asgiref.sync import async_to_sync
from channels.generic.websocket import WebsocketConsumer

from .events import events

class QuestConsumer(WebsocketConsumer):
	"""
	The main consumer for /quest websockets.
	"""
	def __init__(self, *args, **kwargs):
		"""
		Overriden method. Adds dictionary of events and functions to be
		used by self.receive().
		"""
		self.events = events
		super().__init__(*args, **kwargs)

	def connect(self):
		self.quest_id = self.scope['url_route']['kwargs']['quest_id']
		self.group_name = 'quest_' + str(self.quest_id)

		async_to_sync(self.channel_layer.group_add)(
			self.group_name,
			self.channel_name
		)

		self.accept()

	def disconnect(self, close_code):
		async_to_sync(self.channel_layer.group_discard)(
			self.group_name,
			self.channel_name
		)

	def receive(self, text_data):
		"""
		Parses the data received from the WebSocket as JSON and
		dispatches the appropriate event handler.
		"""
		try:
			data = json.loads(text_data)
		except json.JSONDecodeError:
			return
		event = data.get('event')
		if not event or event not in self.events.keys():
			return
		self.events[event](self, data.get('data'))

	def send(self, event, data={}):
		"""
		Overridden method. All data provided (preferably as a dictionary)
		is parsed into JSON before sending to other consumers in the group.
		"""
		data = json.dumps({'event': event, 'data': data})
		async_to_sync(self.channel_layer.group_send)(
			self.group_name,
			{
				'type': 'dispatch_send',
				'message': data
			}
		)
	
	def self_send(self, event, data={}):
		"""
		Like `send`, except only sends events to the paired websocket
		instead of the entire group.
		"""
		data = json.dumps({'event': event, 'data': data})
		self.dispatch_send({'message': data})

	def dispatch_send(self, event):
		"""
		Receives events from other consumers in the group and relays the
		data to the WebSocket.
		"""
		data = event['message']
		super().send(text_data=data)