diff --git a/CHANGELOG.md b/CHANGELOG.md index 992dcff6e..c3e55a2eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). ## Unreleased - Handle `SSLError` exception. (#918) +- Add context to trace functions. (#932) ## 1.0.5 (March 27th, 2024) diff --git a/docs/extensions.md b/docs/extensions.md index 7a24a4181..560dc892a 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -94,7 +94,7 @@ flow of events within `httpcore`. The simplest way to explain this is with an ex ```python import httpcore -def log(event_name, info): +def log(event_name, info, context): print(event_name, info) r = httpcore.request("GET", "https://www.example.com/", extensions={"trace": log}) @@ -120,6 +120,28 @@ The `event_name` and `info` arguments here will be one of the following: * `{event_type}.{event_name}.complete`, `{"return_value": <...>}` * `{event_type}.{event_name}.failed`, `{"exception": <...>}` +The `context` argument here is the dictionary that contains the context of the concrete trace. You can store the data there and access it through the `started`, `complete`, or `failed` stages. + +For example, you can track how much time a particular event took, like so: + +```python +import httpcore +import time + + +def log(event_name, info, context): + _, event_name, stage = event_name.split(".") + if event_name == "start_tls": + if stage == "started": + context["start_time"] = time.monotonic() + elif stage == "complete": + elapsed = time.monotonic() - context["start_time"] + print(f"TLS handshake took {elapsed:.2f} seconds") + + +r = httpcore.request("GET", "https://www.encode.io/", extensions={"trace": log}) +``` + Note that when using the async variant of `httpcore` the handler function passed to `"trace"` must be an `async def ...` function. The following event types are currently exposed... diff --git a/httpcore/_trace.py b/httpcore/_trace.py index b122a53e8..5fe8a5979 100644 --- a/httpcore/_trace.py +++ b/httpcore/_trace.py @@ -24,11 +24,12 @@ def __init__( self.return_value: Any = None self.should_trace = self.debug or self.trace_extension is not None self.prefix = self.logger.name.split(".")[-1] + self.context: Dict[str, Any] = {} def trace(self, name: str, info: Dict[str, Any]) -> None: if self.trace_extension is not None: prefix_and_name = f"{self.prefix}.{name}" - ret = self.trace_extension(prefix_and_name, info) + ret = self.trace_extension(prefix_and_name, info, self.context) if inspect.iscoroutine(ret): # pragma: no cover raise TypeError( "If you are using a synchronous interface, " @@ -67,7 +68,7 @@ def __exit__( async def atrace(self, name: str, info: Dict[str, Any]) -> None: if self.trace_extension is not None: prefix_and_name = f"{self.prefix}.{name}" - coro = self.trace_extension(prefix_and_name, info) + coro = self.trace_extension(prefix_and_name, info, self.context) if not inspect.iscoroutine(coro): # pragma: no cover raise TypeError( "If you're using an asynchronous interface, " diff --git a/tests/_async/test_connection_pool.py b/tests/_async/test_connection_pool.py index 2fc272049..e8fc257e3 100644 --- a/tests/_async/test_connection_pool.py +++ b/tests/_async/test_connection_pool.py @@ -272,7 +272,7 @@ async def test_trace_request(): called = [] - async def trace(name, kwargs): + async def trace(name, kwargs, context): called.append(name) async with httpcore.AsyncConnectionPool(network_backend=network_backend) as pool: @@ -374,7 +374,7 @@ async def test_connection_pool_with_http_exception(): called = [] - async def trace(name, kwargs): + async def trace(name, kwargs, context): called.append(name) async with httpcore.AsyncConnectionPool(network_backend=network_backend) as pool: @@ -427,7 +427,7 @@ async def connect_tcp( called = [] - async def trace(name, kwargs): + async def trace(name, kwargs, context): called.append(name) async with httpcore.AsyncConnectionPool(network_backend=network_backend) as pool: @@ -799,7 +799,7 @@ async def test_http11_upgrade_connection(): called = [] - async def trace(name, kwargs): + async def trace(name, kwargs, context): called.append(name) async with httpcore.AsyncConnectionPool( diff --git a/tests/_sync/test_connection_pool.py b/tests/_sync/test_connection_pool.py index ee303e5cf..e98de6a5e 100644 --- a/tests/_sync/test_connection_pool.py +++ b/tests/_sync/test_connection_pool.py @@ -272,7 +272,7 @@ def test_trace_request(): called = [] - def trace(name, kwargs): + def trace(name, kwargs, context): called.append(name) with httpcore.ConnectionPool(network_backend=network_backend) as pool: @@ -374,7 +374,7 @@ def test_connection_pool_with_http_exception(): called = [] - def trace(name, kwargs): + def trace(name, kwargs, context): called.append(name) with httpcore.ConnectionPool(network_backend=network_backend) as pool: @@ -427,7 +427,7 @@ def connect_tcp( called = [] - def trace(name, kwargs): + def trace(name, kwargs, context): called.append(name) with httpcore.ConnectionPool(network_backend=network_backend) as pool: @@ -799,7 +799,7 @@ def test_http11_upgrade_connection(): called = [] - def trace(name, kwargs): + def trace(name, kwargs, context): called.append(name) with httpcore.ConnectionPool(