Remove Unpack from instrument_flask (#732)

logfire/_internal/integrations/flask.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import Any
 
 from flask.app import Flask
 
@@ -13,31 +13,20 @@
         "    pip install 'logfire[flask]'"
     )
 
-from logfire import Logfire
 from logfire._internal.utils import maybe_capture_server_headers
-
-if TYPE_CHECKING:
-    from wsgiref.types import WSGIEnvironment
-
-    from opentelemetry.trace import Span
-    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
-        excluded_urls: str | None
-        enable_commenter: bool | None
-        commenter_options: dict[str, str] | None
+from logfire.integrations.flask import CommenterOptions, RequestHook, ResponseHook
 
 
 def instrument_flask(
-    logfire_instance: Logfire, app: Flask, capture_headers: bool = False, **kwargs: Unpack[FlaskInstrumentKwargs]
+    app: Flask,
+    *,
+    capture_headers: bool,
+    enable_commenter: bool,
+    commenter_options: CommenterOptions | None,
+    exclude_urls: str | None = None,
+    request_hook: RequestHook | None = None,
+    response_hook: ResponseHook | None = None,
+    **kwargs: Any,
 ):
     """Instrument `app` so that spans are automatically created for each request.
 
@@ -46,9 +35,10 @@ def instrument_flask(
     maybe_capture_server_headers(capture_headers)
     FlaskInstrumentor().instrument_app(  # type: ignore[reportUnknownMemberType]
         app,
-        **{  # type: ignore
-            'tracer_provider': logfire_instance.config.get_tracer_provider(),
-            'meter_provider': logfire_instance.config.get_meter_provider(),
-            **kwargs,
-        },
+        enable_commenter=enable_commenter,
+        commenter_options=commenter_options,
+        excluded_urls=exclude_urls,
+        request_hook=request_hook,
+        response_hook=response_hook,
+        **kwargs,
     )

logfire/_internal/main.py

@@ -82,9 +82,13 @@
     from starlette.websockets import WebSocket
     from typing_extensions import Unpack
 
+    from ..integrations.flask import (
+        CommenterOptions,
+        RequestHook as FlaskRequestHook,
+        ResponseHook as FlaskResponseHook,
+    )
     from .integrations.asgi import ASGIApp, ASGIInstrumentKwargs
     from .integrations.aws_lambda import LambdaEvent, LambdaHandler
-    from .integrations.flask import FlaskInstrumentKwargs
     from .integrations.httpx import AsyncClientKwargs, ClientKwargs, HTTPXInstrumentKwargs
     from .integrations.mysql import MySQLConnection, MySQLInstrumentKwargs
     from .integrations.psycopg import PsycopgInstrumentKwargs
@@ -1393,7 +1397,16 @@ def instrument_psycopg(self, conn_or_module: Any = None, **kwargs: Unpack[Psycop
         return instrument_psycopg(self, conn_or_module, **kwargs)
 
     def instrument_flask(
-        self, app: Flask, *, capture_headers: bool = False, **kwargs: Unpack[FlaskInstrumentKwargs]
+        self,
+        app: Flask,
+        *,
+        capture_headers: bool = False,
+        enable_commenter: bool = True,
+        commenter_options: CommenterOptions | None = None,
+        exclude_urls: str | None = None,
+        request_hook: FlaskRequestHook | None = None,
+        response_hook: FlaskResponseHook | None = None,
+        **kwargs: Any,
     ) -> None:
         """Instrument `app` so that spans are automatically created for each request.
 
@@ -1404,12 +1417,31 @@ def instrument_flask(
         Args:
             app: The Flask app to instrument.
             capture_headers: Set to `True` to capture all request and response headers.
+            enable_commenter: Adds comments to SQL queries performed by Flask, so that database logs have additional context.
+            commenter_options: Configure the tags to be added to the SQL comments.
+                See more about it on the [SQLCommenter Configurations](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/flask/flask.html#sqlcommenter-configurations).
+            exclude_urls: A string containing a comma-delimited list of regexes used to exclude URLs from tracking.
+            request_hook: A function called right after a span is created for a request.
+            response_hook: A function called right before a span is finished for the response.
             **kwargs: Additional keyword arguments to pass to the OpenTelemetry Flask instrumentation.
         """
         from .integrations.flask import instrument_flask
 
         self._warn_if_not_initialized_for_instrumentation()
-        return instrument_flask(self, app, capture_headers=capture_headers, **kwargs)
+        return instrument_flask(
+            app,
+            capture_headers=capture_headers,
+            enable_commenter=enable_commenter,
+            commenter_options=commenter_options,
+            exclude_urls=exclude_urls,
+            request_hook=request_hook,
+            response_hook=response_hook,
+            **{
+                'tracer_provider': self._config.get_tracer_provider(),
+                'meter_provider': self._config.get_meter_provider(),
+                **kwargs,
+            },
+        )
 
     def instrument_starlette(
         self,

logfire/integrations/flask.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, TypedDict
+
+from opentelemetry.trace import Span
+
+if TYPE_CHECKING:
+    from wsgiref.types import WSGIEnvironment
+
+
+RequestHook = Callable[[Span, 'WSGIEnvironment'], None]
+"""A hook that is called before a request is processed."""
+ResponseHook = Callable[[Span, str, 'list[tuple[str, str]]'], None]
+"""A hook that is called after a response is processed."""
+
+
+class CommenterOptions(TypedDict, total=False):
+    """The `commenter_options` parameter for `instrument_flask`."""
+
+    framework: bool
+    """Include the framework name and version in the comment."""
+    route: bool
+    """Include the route name in the comment."""
+    controller: bool
+    """Include the controller name in the comment."""