Support task attr for StreamResponse

aiohttp/test_utils.py

@@ -560,9 +560,11 @@ def make_mocked_request(method, path, headers=None, *,
     time_service.time.return_value = 12345
     time_service.strtime.return_value = "Tue, 15 Nov 1994 08:12:31 GMT"
 
+    task = mock.Mock()
+
     req = Request(message, payload,
                   transport, reader, writer,
-                  time_service,
+                  time_service, task,
                   secure_proxy_ssl_header=secure_proxy_ssl_header)
 
     match_info = UrlMappingMatchInfo({}, mock.Mock())

aiohttp/web.py

@@ -223,7 +223,7 @@ def _make_request(self, message, payload, protocol):
         return web_reqrep.Request(
             message, payload,
             protocol.transport, protocol.reader, protocol.writer,
-            protocol.time_service,
+            protocol.time_service, protocol._request_handler,
             secure_proxy_ssl_header=self._secure_proxy_ssl_header)
 
     @asyncio.coroutine

aiohttp/web_reqrep.py

@@ -51,7 +51,7 @@ class BaseRequest(collections.MutableMapping, HeadersMixin):
                     hdrs.METH_TRACE, hdrs.METH_DELETE}
 
     def __init__(self, message, payload, transport, reader, writer,
-                 time_service, *,
+                 time_service, task, *,
                  secure_proxy_ssl_header=None):
         self._message = message
         self._transport = transport
@@ -69,6 +69,7 @@ def __init__(self, message, payload, transport, reader, writer,
         self._time_service = time_service
         self._state = {}
         self._cache = {}
+        self._task = task
 
     def clone(self, *, method=sentinel, rel_url=sentinel,
               headers=sentinel):
@@ -103,6 +104,7 @@ def clone(self, *, method=sentinel, rel_url=sentinel,
             self._reader,
             self._writer,
             self._time_service,
+            self._task,
             secure_proxy_ssl_header=self._secure_proxy_ssl_header)
 
     # MutableMapping API
@@ -513,6 +515,8 @@ def __init__(self, *, status=200, reason=None, headers=None):
         self._resp_impl = None
         self._eof_sent = False
 
+        self._task = None
+
         if headers is not None:
             self._headers.extend(headers)
         if hdrs.CONTENT_TYPE not in self._headers:
@@ -532,6 +536,10 @@ def started(self):
         warnings.warn('use Response.prepared instead', DeprecationWarning)
         return self.prepared
 
+    @property
+    def task(self):
+        return self._task
+
     @property
     def status(self):
         return self._status
@@ -835,6 +843,7 @@ def _start(self, request):
         resp_impl.headers = headers
 
         self._send_headers(resp_impl)
+        self._task = request._task
         return resp_impl
 
     def _send_headers(self, resp_impl):

aiohttp/web_server.py

@@ -55,12 +55,15 @@ def handle_request(self, message, payload):
         self._request = request
 
         try:
-            resp = yield from self._handler(request)
-        except HTTPException as exc:
-            resp = exc
+            try:
+                resp = yield from self._handler(request)
+            except HTTPException as exc:
+                resp = exc
 
-        resp_msg = yield from resp.prepare(request)
-        yield from resp.write_eof()
+            resp_msg = yield from resp.prepare(request)
+            yield from resp.write_eof()
+        finally:
+            resp._task = None
 
         # notify server about keep-alive
         self.keep_alive(resp.keep_alive)
@@ -121,7 +124,7 @@ def _make_request(self, message, payload, protocol):
         return BaseRequest(
             message, payload,
             protocol.transport, protocol.reader, protocol.writer,
-            protocol.time_service)
+            protocol.time_service, protocol._request_handler)
 
     @asyncio.coroutine
     def shutdown(self, timeout=None):

tests/test_web_functional.py

@@ -1185,3 +1185,26 @@ def handler(request):
     resp = yield from client.get('/')
     assert 200 == resp.status
     assert resp.headers['Date'] == 'Sun, 30 Oct 2016 03:13:52 GMT'
+
+
+@asyncio.coroutine
+def test_response_task(loop, test_client):
+
+    srv_resp = None
+
+    @asyncio.coroutine
+    def handler(request):
+        nonlocal srv_resp
+        srv_resp = web.StreamResponse()
+        assert srv_resp.task is None
+        yield from srv_resp.prepare(request)
+        assert srv_resp.task is not None
+        return srv_resp
+
+    app = web.Application(loop=loop)
+    app.router.add_get('/', handler)
+    client = yield from test_client(app)
+
+    resp = yield from client.get('/')
+    assert 200 == resp.status
+    assert srv_resp.task is None