Add tracking signals for getting request/response bodies.

aiohttp/client.py

@@ -309,7 +309,7 @@ async def _request(self, method, url, *,
                         response_class=self._response_class,
                         proxy=proxy, proxy_auth=proxy_auth, timer=timer,
                         session=self, auto_decompress=self._auto_decompress,
-                        ssl=ssl, proxy_headers=proxy_headers)
+                        ssl=ssl, proxy_headers=proxy_headers, traces=traces)
 
                     # connection timeout
                     try:

aiohttp/client_reqrep.py

@@ -168,7 +168,8 @@ def __init__(self, method, url, *,
                  proxy=None, proxy_auth=None,
                  timer=None, session=None, auto_decompress=True,
                  ssl=None,
-                 proxy_headers=None):
+                 proxy_headers=None,
+                 traces=None):
 
         if loop is None:
             loop = asyncio.get_event_loop()
@@ -209,6 +210,7 @@ def __init__(self, method, url, *,
         if data or self.method not in self.GET_METHODS:
             self.update_transfer_encoding()
         self.update_expect_continue(expect100)
+        self._traces = traces or []
 
     def is_ssl(self):
         return self.url.scheme in ('https', 'wss')
@@ -475,7 +477,10 @@ async def send(self, conn):
             if self.url.raw_query_string:
                 path += '?' + self.url.raw_query_string
 
-        writer = StreamWriter(conn.protocol, conn.transport, self.loop)
+        writer = StreamWriter(
+            conn.protocol, conn.transport, self.loop,
+            on_chunk_sent=self._on_chunk_request_sent
+        )
 
         if self.compress:
             writer.enable_compression(self.compress)
@@ -513,8 +518,9 @@ async def send(self, conn):
             self.method, self.original_url,
             writer=self._writer, continue100=self._continue, timer=self._timer,
             request_info=self.request_info,
-            auto_decompress=self._auto_decompress)
-
+            auto_decompress=self._auto_decompress,
+            traces=self._traces,
+        )
         self.response._post_init(self.loop, self._session)
         return self.response
 
@@ -531,6 +537,10 @@ def terminate(self):
                 self._writer.cancel()
             self._writer = None
 
+    async def _on_chunk_request_sent(self, chunk):
+        for trace in self._traces:
+            await trace.send_request_chunk_sent(chunk)
+
 
 class ClientResponse(HeadersMixin):
 
@@ -555,7 +565,8 @@ class ClientResponse(HeadersMixin):
 
     def __init__(self, method, url, *,
                  writer=None, continue100=None, timer=None,
-                 request_info=None, auto_decompress=True):
+                 request_info=None, auto_decompress=True,
+                 traces=None):
         assert isinstance(url, URL)
 
         self.method = method
@@ -572,6 +583,7 @@ def __init__(self, method, url, *,
         self._timer = timer if timer is not None else TimerNoop()
         self._auto_decompress = auto_decompress
         self._cache = {}  # reqired for @reify method decorator
+        self._traces = traces or []
 
     @property
     def url(self):
@@ -796,6 +808,8 @@ async def read(self):
         if self._content is None:
             try:
                 self._content = await self.content.read()
+                for trace in self._traces:
+                    await trace.send_response_chunk_received(self._content)
             except BaseException:
                 self.close()
                 raise

aiohttp/http_writer.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import collections
+import inspect
 import zlib
 
 from .abc import AbstractStreamWriter
@@ -16,7 +17,12 @@
 
 class StreamWriter(AbstractStreamWriter):
 
-    def __init__(self, protocol, transport, loop):
+    def __init__(self, protocol, transport, loop, on_chunk_sent=None):
+        assert (
+            on_chunk_sent is None or
+            inspect.iscoroutinefunction(on_chunk_sent)
+        )
+
         self._protocol = protocol
         self._transport = transport
 
@@ -30,6 +36,8 @@ def __init__(self, protocol, transport, loop):
         self._compress = None
         self._drain_waiter = None
 
+        self._on_chunk_sent = on_chunk_sent
+
     @property
     def transport(self):
         return self._transport
@@ -62,6 +70,9 @@ async def write(self, chunk, *, drain=True, LIMIT=64*1024):
         writer can't be used after write_eof() method being called.
         write() return drain future.
         """
+        if self._on_chunk_sent is not None:
+            await self._on_chunk_sent(chunk)
+
         if self._compress is not None:
             chunk = self._compress.compress(chunk)
             if not chunk:

aiohttp/tracing.py

@@ -15,7 +15,8 @@
     'TraceConnectionCreateEndParams', 'TraceConnectionReuseconnParams',
     'TraceDnsResolveHostStartParams', 'TraceDnsResolveHostEndParams',
     'TraceDnsCacheHitParams', 'TraceDnsCacheMissParams',
-    'TraceRequestRedirectParams'
+    'TraceRequestRedirectParams',
+    'TraceRequestChunkSentParams', 'TraceResponseChunkReceivedParams',
 )
 
 
@@ -25,6 +26,8 @@ class TraceConfig:
 
     def __init__(self, trace_config_ctx_factory=SimpleNamespace):
         self._on_request_start = Signal(self)
+        self._on_request_chunk_sent = Signal(self)
+        self._on_response_chunk_received = Signal(self)
         self._on_request_end = Signal(self)
         self._on_request_exception = Signal(self)
         self._on_request_redirect = Signal(self)
@@ -47,6 +50,8 @@ def trace_config_ctx(self, trace_request_ctx=None):
 
     def freeze(self):
         self._on_request_start.freeze()
+        self._on_request_chunk_sent.freeze()
+        self._on_response_chunk_received.freeze()
         self._on_request_end.freeze()
         self._on_request_exception.freeze()
         self._on_request_redirect.freeze()
@@ -64,6 +69,14 @@ def freeze(self):
     def on_request_start(self):
         return self._on_request_start
 
+    @property
+    def on_request_chunk_sent(self):
+        return self._on_request_chunk_sent
+
+    @property
+    def on_response_chunk_received(self):
+        return self._on_response_chunk_received
+
     @property
     def on_request_end(self):
         return self._on_request_end
@@ -121,6 +134,18 @@ class TraceRequestStartParams:
     headers = attr.ib(type=CIMultiDict)
 
 
+@attr.s(frozen=True, slots=True)
+class TraceRequestChunkSentParams:
+    """ Parameters sent by the `on_request_chunk_sent` signal"""
+    chunk = attr.ib(type=bytes)
+
+
+@attr.s(frozen=True, slots=True)
+class TraceResponseChunkReceivedParams:
+    """ Parameters sent by the `on_response_chunk_received` signal"""
+    chunk = attr.ib(type=bytes)
+
+
 @attr.s(frozen=True, slots=True)
 class TraceRequestEndParams:
     """ Parameters sent by the `on_request_end` signal"""
@@ -213,6 +238,20 @@ async def send_request_start(self, method, url, headers):
             TraceRequestStartParams(method, url, headers)
         )
 
+    async def send_request_chunk_sent(self, chunk):
+        return await self._trace_config.on_request_chunk_sent.send(
+            self._session,
+            self._trace_config_ctx,
+            TraceRequestChunkSentParams(chunk)
+        )
+
+    async def send_response_chunk_received(self, chunk):
+        return await self._trace_config.on_response_chunk_received.send(
+            self._session,
+            self._trace_config_ctx,
+            TraceResponseChunkReceivedParams(chunk)
+        )
+
     async def send_request_end(self, method, url, headers, response):
         return await self._trace_config.on_request_end.send(
             self._session,

docs/tracing_reference.rst

@@ -34,8 +34,8 @@ Overview
      exception[shape=flowchart.terminator, description="on_request_exception"];
 
      acquire_connection[description="Connection acquiring"];
-     got_response;
-     send_request;
+     got_response[description="on_response_chunk_received"];
+     send_request[description="on_request_chunk_sent"];
 
      start -> acquire_connection;
      acquire_connection -> send_request;
@@ -147,6 +147,20 @@ TraceConfig
 
       ``params`` is :class:`aiohttp.TraceRequestStartParams` instance.
 
+   .. attribute:: on_request_chunk_sent
+
+      Property that gives access to the signals that will be executed
+      when a chunk of request body is sent.
+
+      ``params`` is :class:`aiohttp.TraceRequestChunkSentParams` instance.
+
+   .. attribute:: on_response_chunk_received
+
+      Property that gives access to the signals that will be executed
+      when a chunk of response body is received.
+
+      ``params`` is :class:`aiohttp.TraceResponseChunkReceivedParams` instance.
+
    .. attribute:: on_request_redirect
 
       Property that gives access to the signals that will be executed when a
@@ -259,6 +273,31 @@ TraceRequestStartParams
 
        Headers that will be used for the request, can be mutated.
 
+
+TraceRequestChunkSentParams
+---------------------------
+
+.. class:: TraceRequestChunkSentParams
+
+   See :attr:`TraceConfig.on_request_chunk_sent` for details.
+
+   .. attribute:: chunk
+
+       Bytes of chunk sent
+
+
+TraceResponseChunkSentParams
+----------------------------
+
+.. class:: TraceResponseChunkSentParams
+
+   See :attr:`TraceConfig.on_response_chunk_received` for details.
+
+   .. attribute:: chunk
+
+       Bytes of chunk received
+
+
 TraceRequestEndParams
 ---------------------
 

setup.cfg

@@ -1,5 +1,6 @@
 [pep8]
 max-line-length=79
+ignore=E225,E226
 
 [easy_install]
 zip_ok = false

tests/test_client_session.py

@@ -1,8 +1,10 @@
 import asyncio
 import contextlib
 import gc
+import json
 import re
 from http.cookies import SimpleCookie
+from io import BytesIO
 from unittest import mock
 
 import pytest
@@ -457,33 +459,47 @@ def test_client_session_implicit_loop_warn():
 
 async def test_request_tracing(loop, aiohttp_client):
     async def handler(request):
-        return web.Response()
+        return web.json_response({'ok': True})
 
     app = web.Application()
-    app.router.add_get('/', handler)
+    app.router.add_post('/', handler)
 
     trace_config_ctx = mock.Mock()
     trace_request_ctx = {}
+    body = 'This is request body'
+    gathered_req_body = BytesIO()
+    gathered_res_body = BytesIO()
     on_request_start = mock.Mock(side_effect=asyncio.coroutine(mock.Mock()))
     on_request_redirect = mock.Mock(side_effect=asyncio.coroutine(mock.Mock()))
     on_request_end = mock.Mock(side_effect=asyncio.coroutine(mock.Mock()))
 
+    async def on_request_chunk_sent(session, context, params):
+        gathered_req_body.write(params.chunk)
+
+    async def on_response_chunk_received(session, context, params):
+        gathered_res_body.write(params.chunk)
+
     trace_config = aiohttp.TraceConfig(
         trace_config_ctx_factory=mock.Mock(return_value=trace_config_ctx)
     )
     trace_config.on_request_start.append(on_request_start)
     trace_config.on_request_end.append(on_request_end)
+    trace_config.on_request_chunk_sent.append(on_request_chunk_sent)
+    trace_config.on_response_chunk_received.append(on_response_chunk_received)
     trace_config.on_request_redirect.append(on_request_redirect)
 
     session = await aiohttp_client(app, trace_configs=[trace_config])
 
-    async with session.get('/', trace_request_ctx=trace_request_ctx) as resp:
+    async with session.post(
+            '/', data=body, trace_request_ctx=trace_request_ctx) as resp:
+
+        await resp.json()
 
         on_request_start.assert_called_once_with(
             session.session,
             trace_config_ctx,
             aiohttp.TraceRequestStartParams(
-                hdrs.METH_GET,
+                hdrs.METH_POST,
                 session.make_url('/'),
                 CIMultiDict()
             )
@@ -493,13 +509,16 @@ async def handler(request):
             session.session,
             trace_config_ctx,
             aiohttp.TraceRequestEndParams(
-                hdrs.METH_GET,
+                hdrs.METH_POST,
                 session.make_url('/'),
                 CIMultiDict(),
                 resp
             )
         )
         assert not on_request_redirect.called
+        assert gathered_req_body.getvalue() == body.encode('utf8')
+        assert gathered_res_body.getvalue() == json.dumps(
+            {'ok': True}).encode('utf8')
 
 
 async def test_request_tracing_exception(loop):