__init__.py 2.21 KB
Newer Older
1
import json
echel0n's avatar
echel0n committed
2
from queue import Queue
3

echel0n's avatar
echel0n committed
4
from tornado.websocket import WebSocketHandler
5
6
7
8

import sickrage

clients = set()
echel0n's avatar
echel0n committed
9
message_queue = Queue()
10
11
12
13
14
15
16
17
18
19
20
21


class WebSocketUIHandler(WebSocketHandler):
    """WebSocket handler to send and receive data to and from a web client."""

    def check_origin(self, origin):
        """Allow alternate origins."""
        return True

    def open(self, *args, **kwargs):
        """Client connected to the WebSocket."""
        clients.add(self)
echel0n's avatar
echel0n committed
22
23
        while not message_queue.empty():
            self.write_message(message_queue.get())
24

25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
    def on_message(self, message):
        """Received a message from the client."""
        sickrage.app.log.debug('WebSocket received message from {}: {}'.format(self.request.remote_ip, message))

    def data_received(self, chunk):
        """Received a streamed data chunk from the client."""
        super(WebSocketUIHandler, self).data_received(chunk)

    def on_close(self):
        """Client disconnected from the WebSocket."""
        clients.remove(self)

    def __repr__(self):
        """Client representation."""
        return '<{} Client: {}>'.format(type(self).__name__, self.request.remote_ip)


class WebSocketMessage(object):
    """Represents a WebSocket message."""

echel0n's avatar
echel0n committed
45
    def __init__(self, message_type, data):
46
47
        """
        Construct a new WebSocket message.
echel0n's avatar
echel0n committed
48
        :param message_type: A string representing the type of message (e.g. notification)
49
50
        :param data: A JSON-serializable object containing the message data.
        """
echel0n's avatar
echel0n committed
51
        self.type = message_type
52
53
54
55
56
57
        self.data = data

    @property
    def content(self):
        """Get the message content."""
        return {
echel0n's avatar
echel0n committed
58
            'type': self.type,
59
60
61
            'data': self.data
        }

62
63
64
65
    def json(self):
        """Return the message content as a JSON-serialized string."""
        return json.dumps(self.content)

66
67
    def push(self):
        """Push the message to all connected WebSocket clients."""
68
        # message_queue.put(self.json())
69
        for client in clients.copy():
70
            try:
71
72
                message = self.json()
                sickrage.app.wserver.io_loop.add_callback(client.write_message, message)
73
74
            except AssertionError:
                continue