From: Ben Darnell <ben@bendarnell.com>
Date: Sun, 8 Jan 2012 01:31:56 +0000 (-0800)
Subject: Add an explicit binary option to WebSocketHandler.write_message.
X-Git-Tag: v2.2.0~47
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=33cd3456d5c773d4ee53a2d5ce71ce62f6b4209f;p=thirdparty%2Ftornado.git

Add an explicit binary option to WebSocketHandler.write_message.

Switching automatically based on python's bytes vs unicode types is
error-prone in python2 where e.g. json_encode returns bytes.

Closes #429.
---

diff --git a/maint/test/websocket/server.py b/maint/test/websocket/server.py
index 276343ae1..83e2da619 100644
--- a/maint/test/websocket/server.py
+++ b/maint/test/websocket/server.py
@@ -2,6 +2,7 @@
 
 from tornado.ioloop import IOLoop
 from tornado.options import define, options, parse_command_line
+from tornado.util import bytes_type
 from tornado.websocket import WebSocketHandler
 from tornado.web import Application
 
@@ -9,7 +10,7 @@ define('port', default=9000)
 
 class EchoHandler(WebSocketHandler):
     def on_message(self, message):
-        self.write_message(message)
+        self.write_message(message, binary=isinstance(message, bytes_type))
 
 if __name__ == '__main__':
     parse_command_line()
diff --git a/tornado/websocket.py b/tornado/websocket.py
index ce8e41dd6..4709ea824 100644
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -100,9 +100,11 @@ class WebSocketHandler(tornado.web.RequestHandler):
             self.ws_connection = WebSocketProtocol76(self)
             self.ws_connection.accept_connection()
 
-    def write_message(self, message):
+    def write_message(self, message, binary=False):
         """Sends the given message to the client of this Web Socket."""
-        self.ws_connection.write_message(message)
+        if isinstance(message, dict):
+            message = tornado.escape.json_encode(message)
+        self.ws_connection.write_message(message, binary=binary)
 
     def open(self, *args, **kwargs):
         """Invoked when a new WebSocket is opened."""
@@ -312,10 +314,11 @@ class WebSocketProtocol76(WebSocketProtocol):
         self.client_terminated = True
         self.close()
 
-    def write_message(self, message):
+    def write_message(self, message, binary=False):
         """Sends the given message to the client of this Web Socket."""
-        if isinstance(message, dict):
-            message = tornado.escape.json_encode(message)
+        if binary:
+            raise ValueError(
+                "Binary messages not supported by this version of websockets")
         if isinstance(message, unicode):
             message = message.encode("utf-8")
         assert isinstance(message, bytes_type)
@@ -405,15 +408,13 @@ class WebSocketProtocol8(WebSocketProtocol):
         frame += data
         self.stream.write(frame)
 
-    def write_message(self, message):
+    def write_message(self, message, binary=False):
         """Sends the given message to the client of this Web Socket."""
-        if isinstance(message, dict):
-            message = tornado.escape.json_encode(message)
-        if isinstance(message, unicode):
-            opcode = 0x1
-            message = message.encode("utf-8")
-        else:
+        if binary:
             opcode = 0x2
+        else:
+            opcode = 0x1
+        message = tornado.escape.utf8(message)
         assert isinstance(message, bytes_type)
         self._write_frame(True, opcode, message)