Tornado Websocket router for receiving Redux actions

websocket_handlers.py

def get_websocket_urls(routes):
    """Generate list of all /page/websocket routes from a list of routes"""
    for route in routes:
        url, handler = route[0], route[1]
        if hasattr(handler, 'socket'):
            ws_url = (url if url.endswith('/') else url + '/') + 'websocket'
            yield (ws_url, handler.socket.Handler)


class BaseSocketHandler(BaseHandler, WebSocketHandler):
    def check_origin(self, origin):
        if DEBUG: return True

        domain = origin.split('://')[-1]
        if domain == UI_DOMAIN:
            return True

        logger.warning('WEBSOCKET ORIGIN DOMAIN INCORRECT: {0} instead of {1}'
                       .format(domain, UI_DOMAIN))
        return False

    def on_message(self, message):
        if hasattr(self, 'on_json'):
            self.on_json(json.loads(message))


class RoutedSocketHandler(BaseSocketHandler):
    """Don't use RoutedSocketHandler directly, see SocketRouter docs."""

    ROUTING_KEY = 'type'   # json message field to match against routes
    routes = []

    @staticmethod
    def default_route(self, message):
        logger.error('Unrecognized websocket msg: {0}'.format(message))

    def on_json(self, message):
        """send messages to appropriate handler based on msg[ROUTING_KEY]"""
        action = message.get(self.ROUTING_KEY)

        if not action:
            return self.default_route(message)

        # run through route patterns looking for a match to handle the msg
        for pattern, handler in self.routes:
            # if pattern is a compiled regex, try matching it
            if hasattr(pattern, 'match') and pattern.match(action):
                break

            # if pattern is just as str, check for exact match
            if action == pattern:
                break
        else:
            def handler(_self, message):
                self.default_route(message)

        try:
            return handler(self, message)
        except Exception as e:
            if DEBUG:
                self.write_message({'success': False, 'errors': [str(e)]})
            raise e


class SocketRouter(object):
    """
    Routes websocket messages based on message['type'] to handlers
    defined using the @socket.route('pattern') decorator.

    Usage:
        urls = [
            ('/mypage', MyPage),
            ('/mypage/websocket', MyPage.socket.Handler),  # add this automatically using get_websocket_urls
        ]

        class MyPage(BaseHandler):
            socket = SocketRouter()

            def get(self):
                self.render('my_page_template.html')

            @socket.open
            def handle_open(self):
                # self refers to MyPage.socket.Handler(), not MyPage()
                self.write_message({'connected': True})

            @socket.route(pattern('HELLO.*'))
            def handle_hello(self, data):
                self.write_message({'recvd_hello_data': data})

            @socket.on_json
            def handle_message(self, data):
                # if present, all messages go to this method skipping routing
                self.write_message({'recvd_message': data})
    """

    # TODO: refactor into a metaclass that describes RoutedSocketHandlers

    def __init__(self):
        class SocketHandler(RoutedSocketHandler):
            routes = []  # very important, otherwise all views share all routes
            pass
        self.Handler = SocketHandler

    def route(self, pattern):
        """accept exact strings, or compiled regex patterns"""
        def wrapper(handler):
            self.Handler.routes.append((pattern, handler))
            return handler
        return wrapper

    def set_on_handler(self, attr):
        """decorator to set arbitrary methods directly on the SocketHandler"""
        def wrapper(value):
            setattr(self.Handler, attr, value)
            return value
        return wrapper

    def __getattr__(self, attr):
        """allows you to decorate the handler directly e.g. @socket.on_close"""
        return self.set_on_handler(attr)

    def __setattr__(self, attr, value):
        """allows you to set properties directly on the SocketHandler class"""
        if attr == 'Handler':
            self.__dict__['Handler'] = value
        else:
            setattr(self.Handler, attr, value)