webscoket.py

#!/usr/bin/env python
# -*- coding: utf_8 -*-
import websocket
from flask import json
from tornado.websocket import WebSocketHandler, WebSocketClosedError
from utils import ObjectDict
from .log import logger


class WebsocketResponse(ObjectDict):
    def __init__(self, type="Response", api_code=-1, message='fail', data=None, traceback=""):
        super(WebsocketResponse, self).__init__()
        self.type = type
        self.code = api_code
        self.message = message
        self.data = data
        self.traceback = traceback
        if isinstance(self.data, list):
            self.count = len(self.data)


def pushmessage(func):
    def send(*agrs, **kwargs):
        self = agrs[0]
        ret = func(*agrs, **kwargs)
        if ret:
            msg, binary = ret
            try:
                if isinstance(msg, WebsocketResponse) or isinstance(msg, dict):
                    self.write_message(json.dumps(msg), binary)
                elif isinstance(msg, str) or isinstance(msg, str):
                    self.write_message(msg, binary)
                else:
                    self.write_message(repr(msg), binary)
            except WebSocketClosedError as e:
                logger.error(e)
                self.on_close()
    return send

class BaseWebsocket(WebSocketHandler):

    handlers = {}

    def open(self):
        logger.warn("websocket of %s is opened", repr(self))
        className = self.__class__.__name__
        if className not in self.handlers:
            self.handlers[className] = set()
        self.handlers[className].add(self)

    @pushmessage
    def send(self, message, binary=False):
        return message, binary

    def on_close(self):
        logger.warn("websocket of %s is closed", repr(self))
        className = self.__class__.__name__
        if className in self.handlers:
            self.handlers[className].remove(self)

    def check_origin(self, origin):
        return True

    @classmethod
    def boardcastMessage(cls, message, binary=False):
        className = cls.__name__
        if className in cls.handlers:
            for handler in cls.handlers[className]:
                handler.send(message, binary)


class NotifyHandler(BaseWebsocket):
    """
        建立与web前端的通信连接，发送状态信息报文
    """

    def open(self):
        super(NotifyHandler, self).open()

    def on_message(self, message):
        pass


class ThreadNotifyHandler(BaseWebsocket):
    """
        建立与tornado主线程与子线程之间的通信连接
    """
    def open(self):
        super(ThreadNotifyHandler, self).open()

    def on_message(self, message):
        NotifyHandler.boardcastMessage(message)


class Wsclient(object):
    """
        用于和ThreadNotifyHandler建立websocket连接的客户端
        使用方式：在子线程中建立到达threadnotify路由的websocket客户端如
        wsclient = Wsclient()
        wsclient.send(WebsocketResponse("dashboard", 0, "success", data=cache))
    """
    def __init__(self):
        super(Wsclient, self).__init__()
        self.wsclient = websocket.WebSocket()
        self.wsclient.connect("ws://localhost:7777/ws/api/v1/threadnotify")

    def send(self, msg):
        if isinstance(msg, dict):
            msg = json.dumps(msg)
        self.wsclient.send(msg)