Add close argument to send/group_send

This commit is contained in:
Andrew Godwin 2016-07-19 08:52:39 -04:00
parent 32e047a320
commit e15f6ead6f

View File

@ -87,25 +87,33 @@ class WebsocketConsumer(BaseConsumer):
""" """
pass pass
def send(self, text=None, bytes=None): def send(self, text=None, bytes=None, close=False):
""" """
Sends a reply back down the WebSocket Sends a reply back down the WebSocket
""" """
message = {}
if close:
message["close"] = True
if text is not None: if text is not None:
self.message.reply_channel.send({"text": text}) message["text"] = text
elif bytes is not None: elif bytes is not None:
self.message.reply_channel.send({"bytes": bytes}) message["bytes"] = bytes
else: else:
raise ValueError("You must pass text or bytes") raise ValueError("You must pass text or bytes")
self.message.reply_channel.send(message)
@classmethod @classmethod
def group_send(cls, name, text=None, bytes=None): def group_send(cls, name, text=None, bytes=None, close=False):
message = {}
if close:
message["close"] = True
if text is not None: if text is not None:
Group(name).send({"text": text}) message["text"] = text
elif bytes is not None: elif bytes is not None:
Group(name).send({"bytes": bytes}) message["bytes"] = bytes
else: else:
raise ValueError("You must pass text or bytes") raise ValueError("You must pass text or bytes")
Group(name).send(message)
def close(self): def close(self):
""" """
@ -148,15 +156,15 @@ class JsonWebsocketConsumer(WebsocketConsumer):
""" """
pass pass
def send(self, content): def send(self, content, close=False):
""" """
Encode the given content as JSON and send it to the client. Encode the given content as JSON and send it to the client.
""" """
super(JsonWebsocketConsumer, self).send(text=json.dumps(content)) super(JsonWebsocketConsumer, self).send(text=json.dumps(content), close=close)
@classmethod @classmethod
def group_send(cls, name, content): def group_send(cls, name, content, close=False):
WebsocketConsumer.group_send(name, json.dumps(content)) WebsocketConsumer.group_send(name, json.dumps(content), close=close)
class WebsocketDemultiplexer(JsonWebsocketConsumer): class WebsocketDemultiplexer(JsonWebsocketConsumer):
@ -200,8 +208,11 @@ class WebsocketDemultiplexer(JsonWebsocketConsumer):
self.message.reply_channel.send(self.encode(stream, payload)) self.message.reply_channel.send(self.encode(stream, payload))
@classmethod @classmethod
def group_send(cls, name, stream, payload): def group_send(cls, name, stream, payload, close=False):
Group(name).send(cls.encode(stream, payload)) message = cls.encode(stream, payload)
if close:
message["close"] = True
Group(name).send(message)
@classmethod @classmethod
def encode(cls, stream, payload): def encode(cls, stream, payload):