httpx2 support in logfire.instrument_httpx() (#2095)

Author: alexmojaki

logfire/_internal/cli/run.py

@@ -7,16 +7,17 @@
 from email.headerregistry import ContentTypeHeader
 from email.policy import EmailPolicy
 from functools import cached_property, lru_cache
+from importlib import import_module
+from types import ModuleType
 from typing import TYPE_CHECKING, Any, Literal, cast
 
-import httpx
 from opentelemetry.trace import NonRecordingSpan, Span, use_span
 
 from logfire._internal.config import GLOBAL_CONFIG
 from logfire._internal.stack_info import warn_at_user_stacklevel
 
 try:
-    from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+    from opentelemetry.instrumentation import httpx as otel_httpx
 
     from logfire.integrations.httpx import (
         AsyncRequestHook,
@@ -37,15 +38,83 @@
 from logfire._internal.main import set_user_attributes_on_raw_span
 from logfire._internal.utils import handle_internal_errors
 
+HTTPXClientInstrumentor = otel_httpx.HTTPXClientInstrumentor
+
 if TYPE_CHECKING:
     from typing import ParamSpec
 
+    import httpx
+    import httpx2
+
+    HTTPX2ClientInstrumentor: type[otel_httpx.HTTPX2ClientInstrumentor] | None = otel_httpx.HTTPX2ClientInstrumentor
+
     P = ParamSpec('P')
+else:
+    HTTPX2ClientInstrumentor: Any = getattr(otel_httpx, 'HTTPX2ClientInstrumentor', None)
+
+
+def _import_optional_client_module(name: str) -> ModuleType | None:
+    try:
+        return import_module(name)
+    except ModuleNotFoundError as error:
+        if error.name != name:
+            raise
+        return None
+
+
+_httpx = _import_optional_client_module('httpx')
+_httpx2 = _import_optional_client_module('httpx2')
+
+
+def _client_module_types(name: str) -> tuple[type[Any], ...]:
+    return tuple(cast(type[Any], getattr(module, name)) for module in (_httpx, _httpx2) if module is not None)
+
+
+_HTTPX_BYTE_STREAM_TYPES = _client_module_types('ByteStream')
+_HTTPX_RESPONSE_TYPES = _client_module_types('Response')
+_HTTPX_RESPONSE_NOT_READ_TYPES = cast(tuple[type[BaseException], ...], _client_module_types('ResponseNotRead'))
+
+
+def _httpx2_instrumentation_error() -> RuntimeError:
+    return RuntimeError(
+        'Instrumenting `httpx2` requires `opentelemetry-instrumentation-httpx>=0.65b0`.\n'
+        'You can update this with:\n'
+        "    pip install -U 'logfire[httpx]' 'opentelemetry-instrumentation-httpx>=0.65b0'"
+    )
+
+
+def _instrumentors_for_installed_modules() -> list[Any]:
+    instrumentors: list[Any] = []
+    if _httpx is not None:
+        instrumentors.append(HTTPXClientInstrumentor())
+    if _httpx2 is not None:
+        if HTTPX2ClientInstrumentor is None:
+            if not instrumentors:
+                raise _httpx2_instrumentation_error()
+            warn_at_user_stacklevel(
+                '`httpx2` will not be instrumented because it requires `opentelemetry-instrumentation-httpx>=0.65b0`.',
+                UserWarning,
+            )
+        else:
+            instrumentors.append(HTTPX2ClientInstrumentor())
+    if not instrumentors:
+        raise RuntimeError('`logfire.instrument_httpx()` requires either the `httpx` or `httpx2` package.')
+    return instrumentors
+
+
+def _instrumentor_for_client(client: Any) -> tuple[Any, bool]:
+    if _httpx is not None and isinstance(client, (_httpx.Client, _httpx.AsyncClient)):
+        return HTTPXClientInstrumentor(), isinstance(client, _httpx.AsyncClient)
+    if _httpx2 is not None and isinstance(client, (_httpx2.Client, _httpx2.AsyncClient)):
+        if HTTPX2ClientInstrumentor is None:
+            raise _httpx2_instrumentation_error()
+        return HTTPX2ClientInstrumentor(), isinstance(client, _httpx2.AsyncClient)
+    raise TypeError(f'Expected an `httpx` or `httpx2` client, got {type(client).__name__}.')
 
 
 def instrument_httpx(
     logfire_instance: Logfire,
-    client: httpx.Client | httpx.AsyncClient | None,
+    client: httpx.Client | httpx.AsyncClient | httpx2.Client | httpx2.AsyncClient | None,
     capture_all: bool | None,
     capture_headers: bool,
     capture_request_body: bool,
@@ -56,7 +125,7 @@ def instrument_httpx(
     async_response_hook: AsyncResponseHook | None,
     **kwargs: Any,
 ) -> None:
-    """Instrument the `httpx` module so that spans are automatically created for each request.
+    """Instrument the `httpx` and `httpx2` modules so that spans are automatically created for each request.
 
     See the `Logfire.instrument_httpx` method for details.
     """
@@ -100,7 +169,6 @@ def instrument_httpx(
     }
     del kwargs  # make sure only final_kwargs is used
 
-    instrumentor = HTTPXClientInstrumentor()
     logfire_instance = logfire_instance.with_settings(custom_scope_suffix='httpx')
 
     if client is None:
@@ -125,9 +193,11 @@ def instrument_httpx(
             logfire_instance,
         )
 
-        instrumentor.instrument(**final_kwargs)
+        for instrumentor in _instrumentors_for_installed_modules():
+            instrumentor.instrument(**final_kwargs)
     else:
-        if isinstance(client, httpx.AsyncClient):
+        instrumentor, is_async = _instrumentor_for_client(client)
+        if is_async:
             request_hook = make_async_request_hook(
                 request_hook,
                 should_capture_request_headers,
@@ -166,7 +236,7 @@ def instrument_httpx(
 
 
 class LogfireHttpxInfoMixin:
-    headers: httpx.Headers
+    headers: httpx.Headers | httpx2.Headers
 
     @property
     def content_type_header_object(self) -> ContentTypeHeader:
@@ -214,7 +284,7 @@ def capture_text_as_json(self, attr_name: str, text: str):
 
     @property
     def body_is_streaming(self):
-        return not isinstance(self.stream, httpx.ByteStream)
+        return not isinstance(self.stream, _HTTPX_BYTE_STREAM_TYPES)
 
     @property
     def content_type_charset(self):
@@ -263,13 +333,13 @@ def hook(span: LogfireSpan):
         self.on_response_read(hook)
 
     @cached_property
-    def response(self) -> httpx.Response:
+    def response(self) -> httpx.Response | httpx2.Response:
         frame = inspect.currentframe().f_back.f_back  # pyright: ignore[reportOptionalMemberAccess]
         while frame:  # pragma: no branch
             response = frame.f_locals.get('response')
             frame = frame.f_back
-            if isinstance(response, httpx.Response):
-                return response
+            if isinstance(response, _HTTPX_RESPONSE_TYPES):
+                return cast('httpx.Response | httpx2.Response', response)
         raise RuntimeError('Could not find the response object')  # pragma: no cover
 
     def on_response_read(self, hook: Callable[[LogfireSpan], None]):
@@ -301,7 +371,7 @@ def read() -> bytes:
             try:
                 # Only log the body the first time it's read
                 return response.content
-            except httpx.ResponseNotRead:
+            except _HTTPX_RESPONSE_NOT_READ_TYPES:
                 return hook(original_read)
 
         response.read = read
@@ -315,7 +385,7 @@ async def aread() -> bytes:
             try:
                 # Only log the body the first time it's read
                 return response.content
-            except httpx.ResponseNotRead:
+            except _HTTPX_RESPONSE_NOT_READ_TYPES:
                 return await hook(original_aread)
 
         response.aread = aread
@@ -477,7 +547,7 @@ def run_hook(hook: Callable[P, Any] | None, *args: P.args, **kwargs: P.kwargs) -
 
 
 def capture_request_or_response_headers(
-    span: Span, headers: httpx.Headers, request_or_response: Literal['request', 'response']
+    span: Span, headers: httpx.Headers | httpx2.Headers, request_or_response: Literal['request', 'response']
 ) -> None:
     span.set_attributes(
         {
@@ -489,11 +559,13 @@ def capture_request_or_response_headers(
 
 CODES_FOR_METHODS_WITH_DATA_PARAM = [
     inspect.unwrap(method).__code__
+    for module in (_httpx, _httpx2)
+    if module is not None
     for method in [
-        httpx.Client.request,
-        httpx.Client.stream,
-        httpx.AsyncClient.request,
-        httpx.AsyncClient.stream,
+        module.Client.request,
+        module.Client.stream,
+        module.AsyncClient.request,
+        module.AsyncClient.stream,
     ]
 ]
 

logfire/_internal/integrations/httpx.py

@@ -75,6 +75,7 @@
 
     import anthropic
     import httpx
+    import httpx2
     import openai
     import pydantic_ai.models
     import requests
@@ -1537,7 +1538,7 @@ def instrument_asyncpg(self, capture_parameters: bool = False, **kwargs: Any) ->
     @overload
     def instrument_httpx(
         self,
-        client: httpx.Client,
+        client: httpx.Client | httpx2.Client,
         *,
         capture_all: bool = False,
         capture_headers: bool = False,
@@ -1551,7 +1552,7 @@ def instrument_httpx(
     @overload
     def instrument_httpx(
         self,
-        client: httpx.AsyncClient,
+        client: httpx.AsyncClient | httpx2.AsyncClient,
         *,
         capture_all: bool = False,
         capture_headers: bool = False,
@@ -1580,7 +1581,7 @@ def instrument_httpx(
 
     def instrument_httpx(
         self,
-        client: httpx.Client | httpx.AsyncClient | None = None,
+        client: httpx.Client | httpx.AsyncClient | httpx2.Client | httpx2.AsyncClient | None = None,
         *,
         capture_all: bool | None = None,
         capture_headers: bool = False,
@@ -1592,17 +1593,18 @@ def instrument_httpx(
         async_response_hook: HttpxAsyncResponseHook | None = None,
         **kwargs: Any,
     ) -> None:
-        """Instrument the `httpx` module so that spans are automatically created for each request.
+        """Instrument the `httpx` and `httpx2` modules so that spans are automatically created for each request.
 
-        Optionally, pass an `httpx.Client` instance to instrument only that client.
+        Optionally, pass an `httpx` or `httpx2` client instance to instrument only that client.
 
         Uses the
         [OpenTelemetry HTTPX Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/httpx/httpx.html)
-        library, specifically `HTTPXClientInstrumentor().instrument()`, to which it passes `**kwargs`.
+        library, specifically `HTTPXClientInstrumentor().instrument()` (or `HTTPX2ClientInstrumentor` for `httpx2`),
+        to which it passes `**kwargs`.
 
         Args:
-            client: The `httpx.Client` or `httpx.AsyncClient` instance to instrument.
-                If `None`, the default, all clients will be instrumented.
+            client: The `httpx` or `httpx2` client instance to instrument.
+                If `None`, the default, all clients from both installed libraries will be instrumented.
             capture_all: Set to `True` to capture all HTTP headers, request and response bodies.
                 By default checks the environment variable `LOGFIRE_HTTPX_CAPTURE_ALL`.
             capture_headers: Set to `True` to capture all HTTP headers.

logfire/_internal/main.py

@@ -1,17 +1,21 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable
-from typing import Any, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias
 
-import httpx
 from opentelemetry.trace import Span
 
-# TODO(Marcelo): When https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3098/ gets merged,
-# and the next version of `opentelemetry-instrumentation-httpx` is released, we can just do a reimport:
-# from opentelemetry.instrumentation.httpx import RequestInfo as RequestInfo
-# from opentelemetry.instrumentation.httpx import ResponseInfo as ResponseInfo
-# from opentelemetry.instrumentation.httpx import RequestHook as RequestHook
-# from opentelemetry.instrumentation.httpx import ResponseHook as ResponseHook
+if TYPE_CHECKING:
+    import httpx
+    import httpx2
+
+    _HTTPXURL: TypeAlias = httpx.URL | httpx2.URL
+    _HTTPXHeaders: TypeAlias = httpx.Headers | httpx2.Headers
+    _HTTPXStream: TypeAlias = (
+        httpx.SyncByteStream | httpx.AsyncByteStream | httpx2.SyncByteStream | httpx2.AsyncByteStream
+    )
+else:
+    _HTTPXURL = _HTTPXHeaders = _HTTPXStream = Any
 
 
 class RequestInfo(NamedTuple):
@@ -21,9 +25,9 @@ class RequestInfo(NamedTuple):
     """
 
     method: bytes
-    url: httpx.URL
-    headers: httpx.Headers
-    stream: httpx.SyncByteStream | httpx.AsyncByteStream | None
+    url: _HTTPXURL
+    headers: _HTTPXHeaders
+    stream: _HTTPXStream | None
     extensions: dict[str, Any] | None
 
 
@@ -34,8 +38,8 @@ class ResponseInfo(NamedTuple):
     """
 
     status_code: int
-    headers: httpx.Headers
-    stream: httpx.SyncByteStream | httpx.AsyncByteStream | None
+    headers: _HTTPXHeaders
+    stream: _HTTPXStream | None
     extensions: dict[str, Any] | None
 
 

logfire/integrations/httpx.py

@@ -106,6 +106,7 @@ pytest_logfire = "logfire._internal.integrations.pytest"
 [dependency-groups]
 dev = [
     "httpx >= 0.27.2",
+    "httpx2 >= 2.0.0",
     "aiohttp >= 3.10.9",
     "redis >= 5.1.1",
     "pymongo >= 4.10.1",