Return a Future from websocket write_message methods.

Fixes #1478.

diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
index 23a4324ce66193ef8ee67c9c16afe2fdfd51033c..7b47214df71c03ef168c0e2f132ecac37b908989 100644 (file)
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -130,7 +130,7 @@ class WebSocketTest(WebSocketBaseTestCase):
     @gen_test
     def test_websocket_gen(self):
         ws = yield self.ws_connect('/echo')
-        ws.write_message('hello')
+        yield ws.write_message('hello')
         response = yield ws.read_message()
         self.assertEqual(response, 'hello')
         yield self.close(ws)

diff --git a/tornado/websocket.py b/tornado/websocket.py
index 58262cc39961cb0a256e6e1dd800ec73e6753b76..3b8b9a1fb4c2103b1cafe8953561a6aa480dc96e 100644 (file)
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -208,12 +208,15 @@ class WebSocketHandler(tornado.web.RequestHandler):
         .. versionchanged:: 3.2
            `WebSocketClosedError` was added (previously a closed connection
            would raise an `AttributeError`)
+
+        .. versionchanged:: 4.3
+           Returns a `.Future` which can be used for flow control.
         """
         if self.ws_connection is None:
             raise WebSocketClosedError()
         if isinstance(message, dict):
             message = tornado.escape.json_encode(message)
-        self.ws_connection.write_message(message, binary=binary)
+        return self.ws_connection.write_message(message, binary=binary)
 
     def select_subprotocol(self, subprotocols):
         """Invoked when a new WebSocket requests specific subprotocols.
@@ -671,7 +674,7 @@ class WebSocketProtocol13(WebSocketProtocol):
         frame += data
         self._wire_bytes_out += len(frame)
         try:
-            self.stream.write(frame)
+            return self.stream.write(frame)
         except StreamClosedError:
             self._abort()
 
@@ -688,7 +691,7 @@ class WebSocketProtocol13(WebSocketProtocol):
         if self._compressor:
             message = self._compressor.compress(message)
             flags |= self.RSV1
-        self._write_frame(True, opcode, message, flags=flags)
+        return self._write_frame(True, opcode, message, flags=flags)
 
     def write_ping(self, data):
         """Send ping frame."""
@@ -970,7 +973,7 @@ class WebSocketClientConnection(simple_httpclient._HTTPConnection):
 
     def write_message(self, message, binary=False):
         """Sends a message to the WebSocket server."""
-        self.protocol.write_message(message, binary)
+        return self.protocol.write_message(message, binary)
 
     def read_message(self, callback=None):
         """Reads a message from the WebSocket server.