user_stream_client.py

from autobahn.twisted.websocket import WebSocketClientFactory, WebSocketClientProtocol


class UserStreamClientProtocol(WebSocketClientProtocol):
  def __init__(self):
    super(UserStreamClientProtocol, self).__init__()

  def onOpen(self):
    self.factory.user_stream.send_message = self.sendMessage
    self.factory.user_stream.initialize()

  def onMessage(self, payload, isbinary):
    # explicit decode for Python 3 compatibility
    self.factory.user_stream.on_message(payload.decode('utf8'))

  def onClose(self, wasclean, code, reason):
    if not wasclean:
      self.factory.user_stream.on_error(
        Exception('WebSocket closed with error - %s : %s' % (code, reason))
      )
    else:
      self.factory.user_stream.on_disconnect('WebSocket closed cleanly - %s : %s' % (code, reason))


class UserStreamClientFactory(WebSocketClientFactory):
  protocol = UserStreamClientProtocol

  def __init__(self, user_stream):
    super(UserStreamClientFactory, self).__init__(user_stream.user_stream_url)
    self.user_stream = user_stream