Add context to trace functions

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)
 

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...

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, "

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(

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(