Return Futures from RequestHandler.flush and HTTP1Connection.write{,_headers}

diff --git a/tornado/http1connection.py b/tornado/http1connection.py
index af8dd83a01dbbfdd9c0b1cfdba7ab83229682639..dd0e7f2cd841f8084ceb30938e1e7db37f80336e 100644 (file)
--- a/tornado/http1connection.py
+++ b/tornado/http1connection.py
@@ -184,6 +184,7 @@ class HTTP1Connection(object):
         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 @@ class HTTP1Connection(object):
         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 @@ class HTTP1Connection(object):
 
     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 @@ class HTTP1Connection(object):
             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

diff --git a/tornado/test/simple_httpclient_test.py b/tornado/test/simple_httpclient_test.py
index 4d852bb289dbe3552a2def2b617230b1e58d48e7..529a71a67727809da898c1a65cd12d6bc5d13489 100644 (file)
--- a/tornado/test/simple_httpclient_test.py
+++ b/tornado/test/simple_httpclient_test.py
@@ -352,11 +352,9 @@ class SimpleHTTPClientTestMixin(object):
 
     @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",

diff --git a/tornado/test/web_test.py b/tornado/test/web_test.py
index 4a45b8471f993bd2635856781f4da81f247dbd15..f1e7536b113feea859ff4a529d13b9a3ede6cb98 100644 (file)
--- a/tornado/test/web_test.py
+++ b/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")
 
 

diff --git a/tornado/web.py b/tornado/web.py
index 260974190e9ff4aceb367caa71a7afe5065f53ae..9d7ceffc7bb3c20539a685f74f0eab7c8a7a3a19 100644 (file)
--- a/tornado/web.py
+++ b/tornado/web.py
@@ -790,14 +790,18 @@ class RequestHandler(object):
             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."""