Add better type hints for HTTPX and AsyncPG (#342)

Author: Kludex

logfire/_internal/integrations/asyncpg.py

@@ -1,9 +1,19 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from opentelemetry.instrumentation.asyncpg import AsyncPGInstrumentor
 
+if TYPE_CHECKING:
+    from typing_extensions import TypedDict, Unpack
+
+    class AsyncPGInstrumentKwargs(TypedDict, total=False):
+        skip_dep_check: bool
+
 
-def instrument_asyncpg():
+def instrument_asyncpg(**kwargs: Unpack[AsyncPGInstrumentKwargs]) -> None:
     """Instrument the `asyncpg` module so that spans are automatically created for each query.
 
     See the `Logfire.instrument_asyncpg` method for details.
     """
-    AsyncPGInstrumentor().instrument()  # type: ignore[reportUnknownMemberType]
+    AsyncPGInstrumentor().instrument(**kwargs)  # type: ignore[reportUnknownMemberType]

logfire/_internal/integrations/httpx.py

@@ -1,9 +1,29 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
 
+if TYPE_CHECKING:
+    from typing import Awaitable, Callable, TypedDict, Unpack
+
+    import httpx
+    from opentelemetry.trace import Span
+
+    RequestHook = Callable[[Span, httpx.Request], None]
+    ResponseHook = Callable[[Span, httpx.Request, httpx.Response], None]
+    AsyncRequestHook = Callable[[Span, httpx.Request], Awaitable[None]]
+    AsyncResponseHook = Callable[[Span, httpx.Request, httpx.Response], Awaitable[None]]
+
+    class HTTPXInstrumentKwargs(TypedDict, total=False):
+        request_hook: RequestHook
+        response_hook: ResponseHook
+        async_request_hook: AsyncRequestHook
+        async_response_hook: AsyncResponseHook
+        skip_dep_check: bool
+
 
-def instrument_httpx(**kwargs: Any):
+def instrument_httpx(**kwargs: Unpack[HTTPXInstrumentKwargs]) -> None:
     """Instrument the `httpx` module so that spans are automatically created for each request.
 
     See the `Logfire.instrument_httpx` method for details.

logfire/_internal/main.py

@@ -9,18 +9,7 @@
 from functools import cached_property, partial
 from time import time
 from types import TracebackType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    ContextManager,
-    Iterable,
-    Literal,
-    Sequence,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterable, Literal, Sequence, TypeVar, Union, cast
 
 import opentelemetry.context as context_api
 import opentelemetry.trace as trace_api
@@ -75,8 +64,10 @@
     from starlette.websockets import WebSocket
     from typing_extensions import Unpack
 
+    from .integrations.asyncpg import AsyncPGInstrumentKwargs
     from .integrations.celery import CeleryInstrumentKwargs
     from .integrations.flask import FlaskInstrumentKwargs
+    from .integrations.httpx import HTTPXInstrumentKwargs
     from .integrations.psycopg import PsycopgInstrumentKwargs
     from .integrations.pymongo import PymongoInstrumentKwargs
     from .integrations.redis import RedisInstrumentKwargs
@@ -1033,14 +1024,14 @@ def instrument_anthropic(
             is_async_client,
         )
 
-    def instrument_asyncpg(self):
+    def instrument_asyncpg(self, **kwargs: Unpack[AsyncPGInstrumentKwargs]) -> None:
         """Instrument the `asyncpg` module so that spans are automatically created for each query."""
         from .integrations.asyncpg import instrument_asyncpg
 
         self._warn_if_not_initialized_for_instrumentation()
-        return instrument_asyncpg()
+        return instrument_asyncpg(**kwargs)
 
-    def instrument_httpx(self, **kwargs: Any):
+    def instrument_httpx(self, **kwargs: Unpack[HTTPXInstrumentKwargs]) -> None:
         """Instrument the `httpx` module so that spans are automatically created for each request.
 
         Uses the