From: Ben Darnell <ben@bendarnell.com>
Date: Sat, 26 Apr 2014 23:11:09 +0000 (-0400)
Subject: Return Futures from RequestHandler.flush and HTTP1Connection.write{,_headers}
X-Git-Tag: v4.0.0b1~91^2~7
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=d007a2f9ae0eae6e4259a30934218226dfa92134;p=thirdparty%2Ftornado.git

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

diff --git a/tornado/http1connection.py b/tornado/http1connection.py
index af8dd83a0..dd0e7f2cd 100644
--- 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 4d852bb28..529a71a67 100644
--- 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 4a45b8471..f1e7536b1 100644
--- 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 260974190..9d7ceffc7 100644
--- 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."""