Add type hints to instrument methods (#320)

logfire/_internal/integrations/flask.py

@@ -1,10 +1,34 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from flask.app import Flask
 from opentelemetry.instrumentation.flask import FlaskInstrumentor
 
+if TYPE_CHECKING:
+    from wsgiref.types import WSGIEnvironment
+
+    from opentelemetry.metrics import MeterProvider
+    from opentelemetry.trace import Span, TracerProvider
+    from typing_extensions import Protocol, TypedDict, Unpack
+
+    class RequestHook(Protocol):
+        def __call__(self, span: Span, environment: WSGIEnvironment) -> None: ...
+
+    class ResponseHook(Protocol):
+        def __call__(self, span: Span, status: str, response_headers: list[tuple[str, str]]) -> None: ...
+
+    class FlaskInstrumentKwargs(TypedDict, total=False):
+        request_hook: RequestHook | None
+        response_hook: RequestHook | None
+        tracer_provider: TracerProvider | None
+        excluded_urls: str | None
+        enable_commenter: bool | None
+        commenter_options: dict[str, str] | None
+        meter_provider: MeterProvider | None
+
 
-def instrument_flask(app: Flask, **kwargs: Any):
+def instrument_flask(app: Flask, **kwargs: Unpack[FlaskInstrumentKwargs]):
     """Instrument `app` so that spans are automatically created for each request.
 
     See the `Logfire.instrument_flask` method for details.

logfire/_internal/integrations/psycopg.py

@@ -11,13 +11,25 @@
 if TYPE_CHECKING:  # pragma: no cover
     from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor
     from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
+    from typing_extensions import TypedDict, Unpack
 
     Instrumentor = PsycopgInstrumentor | Psycopg2Instrumentor
 
+    class CommenterOptions(TypedDict, total=False):
+        db_driver: bool
+        db_framework: bool
+        opentelemetry_values: bool
+
+    class PsycopgInstrumentKwargs(TypedDict, total=False):
+        enable_commenter: bool
+        commenter_options: CommenterOptions
+        skip_dep_check: bool
+
+
 PACKAGE_NAMES = ('psycopg', 'psycopg2')
 
 
-def instrument_psycopg(conn_or_module: Any = None, **kwargs: Any):
+def instrument_psycopg(conn_or_module: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]) -> None:
     """Instrument a `psycopg` connection or module so that spans are automatically created for each query.
 
     See the `Logfire.instrument_psycopg` method for details.
@@ -50,7 +62,7 @@ def instrument_psycopg(conn_or_module: Any = None, **kwargs: Any):
     raise ValueError(f"Don't know how to instrument {conn_or_module!r}")
 
 
-def _instrument_psycopg(name: str, conn: Any = None, **kwargs: Any):
+def _instrument_psycopg(name: str, conn: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]) -> None:
     try:
         instrumentor_module = importlib.import_module(f'opentelemetry.instrumentation.{name}')
     except ImportError:
@@ -78,7 +90,7 @@ def _instrument_psycopg(name: str, conn: Any = None, **kwargs: Any):
         instrumentor.instrument_connection(conn)
 
 
-def check_version(name: str, version: str, instrumentor: Instrumentor):
+def check_version(name: str, version: str, instrumentor: Instrumentor) -> bool:
     with contextlib.suppress(Exception):  # it's not worth raising an exception if this fails somehow.
         for dep in instrumentor.instrumentation_dependencies():
             req = Requirement(dep)  # dep is a string like 'psycopg >= 3.1.0'

logfire/_internal/integrations/pymongo.py

@@ -1,9 +1,35 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 from opentelemetry.instrumentation.pymongo import PymongoInstrumentor
 
+if TYPE_CHECKING:
+    from opentelemetry.metrics import MeterProvider
+    from opentelemetry.trace import TracerProvider
+    from pymongo.monitoring import CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent
+    from typing_extensions import Protocol, TypedDict, Unpack
+
+    class RequestHook(Protocol):
+        def __call__(self, span: Any, event: CommandStartedEvent) -> None: ...
+
+    class ResponseHook(Protocol):
+        def __call__(self, span: Any, event: CommandSucceededEvent) -> None: ...
+
+    class FailedHook(Protocol):
+        def __call__(self, span: Any, event: CommandFailedEvent) -> None: ...
+
+    class PymongoInstrumentKwargs(TypedDict, total=False):
+        request_hook: RequestHook | None
+        response_hook: ResponseHook | None
+        failed_hook: FailedHook | None
+        capture_statement: bool | None
+        tracer_provider: TracerProvider | None
+        meter_provider: MeterProvider | None
+        skip_dep_check: bool
+
 
-def instrument_pymongo(**kwargs: Any):
+def instrument_pymongo(**kwargs: Unpack[PymongoInstrumentKwargs]) -> None:
     """Instrument the `pymongo` module so that spans are automatically created for each operation.
 
     See the `Logfire.instrument_pymongo` method for details.

logfire/_internal/integrations/redis.py

@@ -1,9 +1,30 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 from opentelemetry.instrumentation.redis import RedisInstrumentor
 
+if TYPE_CHECKING:
+    from opentelemetry.metrics import MeterProvider
+    from opentelemetry.trace import Span, TracerProvider
+    from redis import Connection
+    from typing_extensions import Protocol, TypedDict, Unpack
+
+    class RequestHook(Protocol):
+        def __call__(self, span: Span, instance: Connection, *args: Any, **kwargs: Any) -> None: ...
+
+    class ResponseHook(Protocol):
+        def __call__(self, span: Span, instance: Connection, response: Any) -> None: ...
+
+    class RedisInstrumentKwargs(TypedDict, total=False):
+        request_hook: RequestHook | None
+        response_hook: ResponseHook | None
+        tracer_provider: TracerProvider | None
+        meter_provider: MeterProvider | None
+        skip_dep_check: bool
+
 
-def instrument_redis(**kwargs: Any):
+def instrument_redis(**kwargs: Unpack[RedisInstrumentKwargs]) -> None:
     """Instrument the `redis` module so that spans are automatically created for each operation.
 
     See the `Logfire.instrument_redis` method for details.

logfire/_internal/integrations/sqlalchemy.py

@@ -1,9 +1,26 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
 
+if TYPE_CHECKING:
+    from sqlalchemy import Engine
+    from typing_extensions import TypedDict, Unpack
+
+    class CommenterOptions(TypedDict, total=False):
+        db_driver: bool
+        db_framework: bool
+        opentelemetry_values: bool
+
+    class SQLAlchemyInstrumentKwargs(TypedDict, total=False):
+        engine: Engine | None
+        enable_commenter: bool | None
+        commenter_options: CommenterOptions | None
+        skip_dep_check: bool
+
 
-def instrument_sqlalchemy(**kwargs: Any):
+def instrument_sqlalchemy(**kwargs: Unpack[SQLAlchemyInstrumentKwargs]) -> None:
     """Instrument the `sqlalchemy` module so that spans are automatically created for each query.
 
     See the `Logfire.instrument_sqlalchemy` method for details.

logfire/_internal/integrations/starlette.py

@@ -1,10 +1,33 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 from opentelemetry.instrumentation.starlette import StarletteInstrumentor
 from starlette.applications import Starlette
 
+if TYPE_CHECKING:
+    from opentelemetry.metrics import MeterProvider
+    from opentelemetry.trace import Span, TracerProvider
+    from typing_extensions import Protocol, TypedDict, Unpack
+
+    class ServerRequestHook(Protocol):
+        def __call__(self, span: Span, scope: dict[str, Any]): ...
+
+    class ClientRequestHook(Protocol):
+        def __call__(self, span: Span, scope: dict[str, Any]): ...
+
+    class ClientResponseHook(Protocol):
+        def __call__(self, span: Span, message: dict[str, Any]): ...
+
+    class StarletteInstrumentKwargs(TypedDict, total=False):
+        server_request_hook: ServerRequestHook | None
+        client_request_hook: ClientRequestHook | None
+        client_response_hook: ClientResponseHook | None
+        tracer_provider: TracerProvider | None
+        meter_provider: MeterProvider | None
+
 
-def instrument_starlette(app: Starlette, **kwargs: Any):
+def instrument_starlette(app: Starlette, **kwargs: Unpack[StarletteInstrumentKwargs]):
     """Instrument `app` so that spans are automatically created for each request.
 
     See the `Logfire.instrument_starlette` method for details.

logfire/_internal/main.py

@@ -73,7 +73,14 @@
     from starlette.applications import Starlette
     from starlette.requests import Request
     from starlette.websockets import WebSocket
+    from typing_extensions import Unpack
 
+    from .integrations.flask import FlaskInstrumentKwargs
+    from .integrations.psycopg import PsycopgInstrumentKwargs
+    from .integrations.pymongo import PymongoInstrumentKwargs
+    from .integrations.redis import RedisInstrumentKwargs
+    from .integrations.sqlalchemy import SQLAlchemyInstrumentKwargs
+    from .integrations.starlette import StarletteInstrumentKwargs
 
 try:
     from pydantic import ValidationError
@@ -1103,7 +1110,7 @@ def instrument_requests(self, excluded_urls: str | None = None, **kwargs: Any):
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_requests(excluded_urls=excluded_urls, **kwargs)
 
-    def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Any):
+    def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Unpack[PsycopgInstrumentKwargs]) -> None:
         """Instrument a `psycopg` connection or module so that spans are automatically created for each query.
 
         Uses the OpenTelemetry instrumentation libraries for
@@ -1127,7 +1134,7 @@ def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Any):
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_psycopg(conn_or_module, **kwargs)
 
-    def instrument_flask(self, app: Flask, **kwargs: Any):
+    def instrument_flask(self, app: Flask, **kwargs: Unpack[FlaskInstrumentKwargs]) -> None:
         """Instrument `app` so that spans are automatically created for each request.
 
         Uses the
@@ -1139,7 +1146,7 @@ def instrument_flask(self, app: Flask, **kwargs: Any):
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_flask(app, **kwargs)
 
-    def instrument_starlette(self, app: Starlette, **kwargs: Any):
+    def instrument_starlette(self, app: Starlette, **kwargs: Unpack[StarletteInstrumentKwargs]) -> None:
         """Instrument `app` so that spans are automatically created for each request.
 
         Uses the
@@ -1163,7 +1170,7 @@ def instrument_aiohttp_client(self, **kwargs: Any):
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_aiohttp_client(**kwargs)
 
-    def instrument_sqlalchemy(self, **kwargs: Any):
+    def instrument_sqlalchemy(self, **kwargs: Unpack[SQLAlchemyInstrumentKwargs]) -> None:
         """Instrument the `sqlalchemy` module so that spans are automatically created for each query.
 
         Uses the
@@ -1175,7 +1182,7 @@ def instrument_sqlalchemy(self, **kwargs: Any):
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_sqlalchemy(**kwargs)
 
-    def instrument_pymongo(self, **kwargs: Any):
+    def instrument_pymongo(self, **kwargs: Unpack[PymongoInstrumentKwargs]) -> None:
         """Instrument the `pymongo` module so that spans are automatically created for each operation.
 
         Uses the
@@ -1187,7 +1194,7 @@ def instrument_pymongo(self, **kwargs: Any):
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_pymongo(**kwargs)
 
-    def instrument_redis(self, **kwargs: Any):
+    def instrument_redis(self, **kwargs: Unpack[RedisInstrumentKwargs]) -> None:
         """Instrument the `redis` module so that spans are automatically created for each operation.
 
         Uses the