Improve logfire inspect

logfire/_internal/cli.py

@@ -5,13 +5,16 @@
 import argparse
 import functools
 import importlib
+import importlib.metadata
 import importlib.util
 import logging
 import os
 import platform
+import re
 import sys
 import warnings
 import webbrowser
+from collections import defaultdict
 from pathlib import Path
 from typing import Any, Iterator, cast
 from urllib.parse import urljoin, urlparse
@@ -195,6 +198,7 @@ def reader() -> Iterator[bytes]:
     'urllib3',
 }
 OTEL_PACKAGE_LINK = {'aiohttp': 'aiohttp-client', 'tortoise_orm': 'tortoiseorm', 'scikit-learn': 'sklearn'}
+STANDARD_LIBRARY_PACKAGES = {'urllib', 'sqlite3'}
 
 
 def parse_inspect(args: argparse.Namespace) -> None:
@@ -207,11 +211,23 @@ def parse_inspect(args: argparse.Namespace) -> None:
     # Ignore warnings from packages that we don't control.
     warnings.simplefilter('ignore', category=UserWarning)
 
+    packages_to_inspect = OTEL_PACKAGES
+    if args.ignore_standard_library:
+        packages_to_inspect -= STANDARD_LIBRARY_PACKAGES
+    if args.ignore_package:
+        packages_to_inspect -= set(args.ignore_package)
+
+    required_by = _build_required_by()
+
     packages: dict[str, str] = {}
-    for name in OTEL_PACKAGES:
+    for name in packages_to_inspect:
         # Check if the package can be imported (without actually importing it).
         if importlib.util.find_spec(name) is None:
             continue
+        else:
+            # Skip packages that are not required by any other package.
+            if len([req for req in required_by[name] if not req.startswith('opentelemetry')]) == 0:
+                continue
 
         otel_package = OTEL_PACKAGE_LINK.get(name, name)
         otel_package_import = f'opentelemetry.instrumentation.{otel_package}'
@@ -220,15 +236,20 @@ def parse_inspect(args: argparse.Namespace) -> None:
             packages[name] = otel_package
 
     # Drop packages that are dependencies of other packages.
-    if packages.get('starlette') and packages.get('fastapi'):
+    if packages.get('starlette') and len(required_by['starlette'] - {'fastapi'}) == 0:
         del packages['starlette']
+    if packages.get('urllib3') and len(required_by['urllib3'] - {'requests'}) == 0:
+        del packages['urllib3']
 
     for name, otel_package in sorted(packages.items()):
         package_name = otel_package.replace('.', '-')
         import_name = otel_package.replace('-', '_')
         link = f'[link={BASE_OTEL_INTEGRATION_URL}/{import_name}/{import_name}.html]opentelemetry-instrumentation-{package_name}[/link]'
         table.add_row(name, link)
 
+    if not packages:
+        return
+
     console.print(
         'The following packages from your environment have an OpenTelemetry instrumentation that is not installed:'
     )
@@ -244,6 +265,32 @@ def parse_inspect(args: argparse.Namespace) -> None:
         console.print('\n[bold blue]For further information, visit[/bold blue]', end=' ')
         console.print(f'[link={INTEGRATIONS_DOCS_URL}]{INTEGRATIONS_DOCS_URL}[/link]')
 
+    if args.explain:
+        console.print('\n[bold]Explanation:[/]')
+        for name, otel_package in packages.items():
+            required_by_name = required_by[name]
+            if required_by_name:
+                console.print(f'\n[bold]{name}[/] is required by:')
+                for req in required_by_name:
+                    console.print(f'  - {req}')
+
+
+# Package name should only include digits, letters, underscores, and hyphens.
+ONLY_PACKAGE_NAME = re.compile(r'^(?P<name>[a-zA-Z0-9_-]+)')
+
+
+def _build_required_by() -> dict[str, set[str]]:
+    """Build a dictionary of packages and their dependencies."""
+    required_by: defaultdict[str, set[str]] = defaultdict(set)
+    for package in importlib.metadata.distributions():
+        if package.requires is None:
+            continue
+        for requirement in package.requires:
+            match = ONLY_PACKAGE_NAME.match(requirement)
+            assert match is not None
+            required_by[match.group('name')].add(package.metadata['Name'])
+    return required_by
+
 
 def parse_auth(args: argparse.Namespace) -> None:
     """Authenticate with Logfire.
@@ -438,6 +485,9 @@ def _main(args: list[str] | None = None) -> None:
 
     cmd_inspect = subparsers.add_parser('inspect', help=parse_inspect.__doc__)
     cmd_inspect.set_defaults(func=parse_inspect)
+    cmd_inspect.add_argument('--ignore-standard-library', action='store_true', help='ignore standard library packages')
+    cmd_inspect.add_argument('--ignore-package', action='append', help='ignore a package')
+    cmd_inspect.add_argument('--explain', action='store_true', help='explain the output')
 
     cmd_whoami = subparsers.add_parser('whoami', help=parse_whoami.__doc__)
     cmd_whoami.set_defaults(func=parse_whoami)

logfire/_internal/integrations/httpx.py

@@ -71,6 +71,7 @@ def instrument_httpx(
     capture_request_json_body: bool,
     capture_request_text_body: bool,
     capture_response_json_body: bool,
+    capture_response_text_body: bool,
     capture_request_form_data: bool,
     **kwargs: Any,
 ) -> None:
@@ -116,7 +117,11 @@ def instrument_httpx(
             capture_request_form_data,
         )
         final_kwargs['response_hook'] = make_response_hook(
-            response_hook, should_capture_response_headers, capture_response_json_body, logfire_instance
+            response_hook,
+            should_capture_response_headers,
+            capture_response_json_body,
+            capture_response_text_body,
+            logfire_instance,
         )
         final_kwargs['async_request_hook'] = make_async_request_hook(
             async_request_hook,
@@ -126,7 +131,11 @@ def instrument_httpx(
             capture_request_form_data,
         )
         final_kwargs['async_response_hook'] = make_async_response_hook(
-            async_response_hook, should_capture_response_headers, capture_response_json_body, logfire_instance
+            async_response_hook,
+            should_capture_response_headers,
+            capture_response_json_body,
+            capture_response_text_body,
+            logfire_instance,
         )
 
         instrumentor.instrument(**final_kwargs)
@@ -143,7 +152,11 @@ def instrument_httpx(
                 capture_request_form_data,
             )
             response_hook = make_async_response_hook(
-                response_hook, should_capture_response_headers, capture_response_json_body, logfire_instance
+                response_hook,
+                should_capture_response_headers,
+                capture_response_json_body,
+                capture_response_text_body,
+                logfire_instance,
             )
         else:
             request_hook = cast('RequestHook | None', final_kwargs.get('request_hook'))
@@ -157,7 +170,11 @@ def instrument_httpx(
                 capture_request_form_data,
             )
             response_hook = make_response_hook(
-                response_hook, should_capture_response_headers, capture_response_json_body, logfire_instance
+                response_hook,
+                should_capture_response_headers,
+                capture_response_json_body,
+                capture_response_text_body,
+                logfire_instance,
             )
 
         tracer_provider = final_kwargs['tracer_provider']
@@ -310,7 +327,11 @@ def capture_request(
 
 
 def make_response_hook(
-    hook: ResponseHook | None, should_capture_headers: bool, should_capture_json: bool, logfire_instance: Logfire
+    hook: ResponseHook | None,
+    should_capture_headers: bool,
+    should_capture_json: bool,
+    should_capture_text: bool,
+    logfire_instance: Logfire,
 ) -> ResponseHook | None:
     if not should_capture_headers and not should_capture_json and not hook:
         return None
@@ -321,6 +342,8 @@ def new_hook(span: Span, request: RequestInfo, response: ResponseInfo) -> None:
                 capture_response_headers(span, response)
             if should_capture_json:
                 capture_response_json(logfire_instance, response, False)
+            if should_capture_text:
+                capture_response_text(logfire_instance, response, False)
             run_hook(hook, span, request, response)
 
     return new_hook
@@ -330,6 +353,7 @@ def make_async_response_hook(
     hook: ResponseHook | AsyncResponseHook | None,
     should_capture_headers: bool,
     should_capture_json: bool,
+    should_capture_text: bool,
     logfire_instance: Logfire,
 ) -> AsyncResponseHook | None:
     if not should_capture_headers and not should_capture_json and not hook:
@@ -341,6 +365,8 @@ async def new_hook(span: Span, request: RequestInfo, response: ResponseInfo) ->
                 capture_response_headers(span, response)
             if should_capture_json:
                 capture_response_json(logfire_instance, response, True)
+            if should_capture_text:
+                capture_response_text(logfire_instance, response, True)
             await run_async_hook(hook, span, request, response)
 
     return new_hook
@@ -398,6 +424,54 @@ def read(*args: Any, **kwargs: Any):
         response.read = read
 
 
+def capture_response_text(logfire_instance: Logfire, response_info: ResponseInfo, is_async: bool) -> None:
+    if not is_text_type(response_info.headers.get('content-type', '')):
+        return
+
+    frame = inspect.currentframe().f_back.f_back  # type: ignore
+    while frame:
+        response = frame.f_locals.get('response')
+        frame = frame.f_back
+        if isinstance(response, httpx.Response):  # pragma: no branch
+            break
+    else:  # pragma: no cover
+        return
+
+    ctx = get_context()
+    attr_name = 'http.response.body.text'
+
+    if is_async:  # these two branches should be kept almost identical
+        original_aread = response.aread
+
+        async def aread(*args: Any, **kwargs: Any):
+            try:
+                # Only log the body the first time it's read
+                return response.content
+            except httpx.ResponseNotRead:
+                pass
+            with attach_context(ctx), logfire_instance.span('Reading response body') as span:
+                content = await original_aread(*args, **kwargs)
+                span.set_attribute(attr_name, response.text)
+            return content
+
+        response.aread = aread
+    else:
+        original_read = response.read
+
+        def read(*args: Any, **kwargs: Any):
+            try:
+                # Only log the body the first time it's read
+                return response.content
+            except httpx.ResponseNotRead:
+                pass
+            with attach_context(ctx), logfire_instance.span('Reading response body') as span:
+                content = original_read(*args, **kwargs)
+                span.set_attribute(attr_name, response.text)
+            return content
+
+        response.read = read
+
+
 async def run_async_hook(hook: Callable[P, Any] | None, *args: P.args, **kwargs: P.kwargs) -> None:
     if hook:
         result = hook(*args, **kwargs)

tests/test_cli.py

@@ -17,7 +17,7 @@
 
 import logfire._internal.cli
 from logfire import VERSION
-from logfire._internal.cli import OTEL_PACKAGES, main
+from logfire._internal.cli import main
 from logfire._internal.config import LogfireCredentials, sanitize_project_name
 from logfire.exceptions import LogfireConfigError
 
@@ -206,20 +206,6 @@ def test_inspect(
     assert capsys.readouterr().err.startswith('The following packages')
 
 
-def test_inspect_drop_dependant_packages(
-    tmp_dir_cwd: Path, logfire_credentials: LogfireCredentials, capsys: pytest.CaptureFixture[str]
-) -> None:
-    logfire_credentials.write_creds_file(tmp_dir_cwd / '.logfire')
-    with ExitStack() as stack:
-        find_spec = stack.enter_context(patch('importlib.util.find_spec'))
-        find_spec.side_effect = [True, None] * len(OTEL_PACKAGES)
-
-        main(['inspect'])
-        output = capsys.readouterr().err
-        assert 'opentelemetry-instrumentation-fastapi' in output
-        assert 'opentelemetry-instrumentation-starlette' not in output
-
-
 @pytest.mark.parametrize('webbrowser_error', [False, True])
 def test_auth(tmp_path: Path, webbrowser_error: bool) -> None:
     auth_file = tmp_path / 'default.toml'