Return Futures from RequestHandler.flush and HTTP1Connection.write{,_…

…headers}

tornado/http1connection.py

@@ -184,6 +184,7 @@ def _clear_callbacks(self):
         quickly in CPython by breaking up reference cycles.
         """
         self._write_callback = None
+        self._write_future = None
         self._close_callback = None
 
     def set_close_callback(self, callback):
@@ -267,12 +268,19 @@ def write_headers(self, start_line, headers, chunk=None, callback=None,
         for line in lines:
             if b'\n' in line:
                 raise ValueError('Newline in header: ' + repr(line))
-        if not self.stream.closed():
-            self._write_callback = stack_context.wrap(callback)
+        if self.stream.closed():
+            self._write_future = Future()
+            self._write_future.set_exception(iostream.StreamClosedError())
+        else:
+            if callback is not None:
+                self._write_callback = stack_context.wrap(callback)
+            else:
+                self._write_future = Future()
             data = b"\r\n".join(lines) + b"\r\n\r\n"
             if chunk:
                 data += self._format_chunk(chunk)
             self.stream.write(data, self._on_write_complete)
+        return self._write_future
 
     def _format_chunk(self, chunk):
         if self._expected_content_remaining is not None:
@@ -291,10 +299,17 @@ def _format_chunk(self, chunk):
 
     def write(self, chunk, callback=None):
         """Writes a chunk of output to the stream."""
-        if not self.stream.closed():
-            self._write_callback = stack_context.wrap(callback)
+        if self.stream.closed():
+            self._write_future = Future()
+            self._write_future.set_exception(iostream.StreamClosedError())
+        else:
+            if callback is not None:
+                self._write_callback = stack_context.wrap(callback)
+            else:
+                self._write_future = Future()
             self.stream.write(self._format_chunk(chunk),
                               self._on_write_complete)
+        return self._write_future
 
     def finish(self):
         """Finishes the request."""
@@ -327,6 +342,10 @@ def _on_write_complete(self):
             callback = self._write_callback
             self._write_callback = None
             callback()
+        if self._write_future is not None:
+            future = self._write_future
+            self._write_future = None
+            future.set_result(None)
         # _on_write_complete is enqueued on the IOLoop whenever the
         # IOStream's write buffer becomes empty, but it's possible for
         # another callback that runs on the IOLoop before it to

tornado/test/simple_httpclient_test.py

@@ -352,11 +352,9 @@ def sync_body_producer(self, write):
 
     @gen.coroutine
     def async_body_producer(self, write):
-        # TODO: write should return a Future.
-        # wrap it in simple_httpclient or change http1connection?
-        yield gen.Task(write, b'1234')
+        yield write(b'1234')
         yield gen.Task(IOLoop.current().add_callback)
-        yield gen.Task(write, b'5678')
+        yield write(b'5678')
 
     def test_sync_body_producer_chunked(self):
         response = self.fetch("/echo_post", method="POST",

tornado/test/web_test.py

@@ -469,12 +469,13 @@ class EmptyFlushCallbackHandler(RequestHandler):
     @asynchronous
     def get(self):
         # Ensure that the flush callback is run whether or not there
-        # was any output.
+        # was any output.  The gen.Task and direct yield forms are
+        # equivalent.
         yield gen.Task(self.flush)  # "empty" flush, but writes headers
         yield gen.Task(self.flush)  # empty flush
         self.write("o")
-        yield gen.Task(self.flush)  # flushes the "o"
-        yield gen.Task(self.flush)  # empty flush
+        yield self.flush()  # flushes the "o"
+        yield self.flush()  # empty flush
         self.finish("k")
 
 

tornado/web.py

@@ -790,14 +790,18 @@ def flush(self, include_footers=False, callback=None):
             start_line = httputil.ResponseStartLine(self.request.version,
                                                     self._status_code,
                                                     self._reason)
-            self.request.connection.write_headers(start_line, self._headers,
-                                                  chunk, callback=callback)
+            return self.request.connection.write_headers(
+                start_line, self._headers, chunk, callback=callback)
         else:
             for transform in self._transforms:
                 chunk = transform.transform_chunk(chunk, include_footers)
             # Ignore the chunk and only write the headers for HEAD requests
             if self.request.method != "HEAD":
-                self.request.connection.write(chunk, callback=callback)
+                return self.request.connection.write(chunk, callback=callback)
+            else:
+                future = Future()
+                future.set_result(None)
+                return future
 
     def finish(self, chunk=None):
         """Finishes this response, ending the HTTP request."""