Merge branch 'main' into alex/scrub-once

logfire-api/logfire_api/_internal/integrations/celery.pyi

@@ -1,10 +1,6 @@
-from logfire import Logfire as Logfire
-from typing_extensions import TypedDict, Unpack
+from typing import Any
 
-class CeleryInstrumentKwargs(TypedDict, total=False):
-    skip_dep_check: bool
-
-def instrument_celery(logfire_instance: Logfire, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None:
+def instrument_celery(**kwargs: Any) -> None:
     """Instrument the `celery` module so that spans are automatically created for each task.
 
     See the `Logfire.instrument_celery` method for details.

logfire-api/logfire_api/_internal/integrations/requests.pyi

@@ -1,7 +1,8 @@
-from logfire import Logfire as Logfire
-from typing import Any
+import requests
+from opentelemetry.sdk.trace import Span as Span
+from typing import Any, Callable
 
-def instrument_requests(logfire_instance: Logfire, excluded_urls: str | None = None, **kwargs: Any):
+def instrument_requests(excluded_urls: str | None = None, request_hook: Callable[[Span, requests.PreparedRequest], None] | None = None, response_hook: Callable[[Span, requests.PreparedRequest, requests.Response], None] | None = None, **kwargs: Any) -> None:
     """Instrument the `requests` module so that spans are automatically created for each request.
 
     See the `Logfire.instrument_requests` method for details.

logfire-api/logfire_api/_internal/main.pyi

@@ -2,6 +2,7 @@ import anthropic
 import httpx
 import openai
 import opentelemetry.trace as trace_api
+import requests
 from . import async_ as async_
 from ..version import VERSION as VERSION
 from .auto_trace import AutoTraceModule as AutoTraceModule, install_auto_tracing as install_auto_tracing
@@ -13,7 +14,6 @@ from .instrument import instrument as instrument
 from .integrations.asgi import ASGIApp as ASGIApp, ASGIInstrumentKwargs as ASGIInstrumentKwargs
 from .integrations.asyncpg import AsyncPGInstrumentKwargs as AsyncPGInstrumentKwargs
 from .integrations.aws_lambda import AwsLambdaInstrumentKwargs as AwsLambdaInstrumentKwargs, LambdaHandler as LambdaHandler
-from .integrations.celery import CeleryInstrumentKwargs as CeleryInstrumentKwargs
 from .integrations.flask import FlaskInstrumentKwargs as FlaskInstrumentKwargs
 from .integrations.httpx import AsyncClientKwargs as AsyncClientKwargs, ClientKwargs as ClientKwargs, HTTPXInstrumentKwargs as HTTPXInstrumentKwargs
 from .integrations.mysql import MySQLConnection as MySQLConnection, MySQLInstrumentKwargs as MySQLInstrumentKwargs
@@ -556,12 +556,15 @@ class Logfire:
     def instrument_httpx(self, client: httpx.AsyncClient, capture_request_headers: bool = False, capture_response_headers: bool = False, **kwargs: Unpack[AsyncClientKwargs]) -> None: ...
     @overload
     def instrument_httpx(self, client: None = None, capture_request_headers: bool = False, capture_response_headers: bool = False, **kwargs: Unpack[HTTPXInstrumentKwargs]) -> None: ...
-    def instrument_celery(self, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None:
+    def instrument_celery(self, **kwargs: Any) -> None:
         """Instrument `celery` so that spans are automatically created for each task.
 
         Uses the
         [OpenTelemetry Celery Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/celery/celery.html)
         library.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` method, for future compatibility.
         """
     def instrument_django(self, capture_headers: bool = False, is_sql_commentor_enabled: bool | None = None, request_hook: Callable[[Span, HttpRequest], None] | None = None, response_hook: Callable[[Span, HttpRequest, HttpResponse], None] | None = None, excluded_urls: str | None = None, **kwargs: Any) -> None:
         """Instrument `django` so that spans are automatically created for each web request.
@@ -594,13 +597,14 @@ class Logfire:
                 for future compatibility.
 
         """
-    def instrument_requests(self, excluded_urls: str | None = None, **kwargs: Any) -> None:
+    def instrument_requests(self, excluded_urls: str | None = None, request_hook: Callable[[Span, requests.PreparedRequest], None] | None = None, response_hook: Callable[[Span, requests.PreparedRequest, requests.Response], None] | None = None, **kwargs: Any) -> None:
         """Instrument the `requests` module so that spans are automatically created for each request.
 
         Args:
             excluded_urls: A string containing a comma-delimited list of regexes used to exclude URLs from tracking
-            **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods,
-                particularly `request_hook` and `response_hook`.
+            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 `instrument` methods, for future compatibility.
         """
     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.

logfire/_internal/integrations/celery.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import Any
 
 try:
     from opentelemetry.instrumentation.celery import CeleryInstrumentor
@@ -11,24 +11,10 @@
         "    pip install 'logfire[celery]'"
     )
 
-from logfire import Logfire
 
-if TYPE_CHECKING:
-    from typing_extensions import TypedDict, Unpack
-
-    class CeleryInstrumentKwargs(TypedDict, total=False):
-        skip_dep_check: bool
-
-
-def instrument_celery(logfire_instance: Logfire, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None:
+def instrument_celery(**kwargs: Any) -> None:
     """Instrument the `celery` module so that spans are automatically created for each task.
 
     See the `Logfire.instrument_celery` method for details.
     """
-    return CeleryInstrumentor().instrument(
-        **{
-            'tracer_provider': logfire_instance.config.get_tracer_provider(),
-            'meter_provider': logfire_instance.config.get_meter_provider(),
-            **kwargs,
-        }
-    )
+    return CeleryInstrumentor().instrument(**kwargs)

logfire/_internal/integrations/requests.py

@@ -1,4 +1,9 @@
-from typing import Any, Optional
+from __future__ import annotations
+
+from typing import Any, Callable
+
+import requests
+from opentelemetry.sdk.trace import Span
 
 try:
     from opentelemetry.instrumentation.requests import RequestsInstrumentor
@@ -9,19 +14,20 @@
         "    pip install 'logfire[requests]'"
     )
 
-from logfire import Logfire
-
 
-def instrument_requests(logfire_instance: Logfire, excluded_urls: Optional[str] = None, **kwargs: Any):
+def instrument_requests(
+    excluded_urls: str | None = None,
+    request_hook: Callable[[Span, requests.PreparedRequest], None] | None = None,
+    response_hook: Callable[[Span, requests.PreparedRequest, requests.Response], None] | None = None,
+    **kwargs: Any,
+) -> None:
     """Instrument the `requests` module so that spans are automatically created for each request.
 
     See the `Logfire.instrument_requests` method for details.
     """
     RequestsInstrumentor().instrument(
         excluded_urls=excluded_urls,
-        **{
-            'tracer_provider': logfire_instance.config.get_tracer_provider(),
-            'meter_provider': logfire_instance.config.get_meter_provider(),
-            **kwargs,
-        },
+        request_hook=request_hook,
+        response_hook=response_hook,
+        **kwargs,
     )

logfire/_internal/main.py

@@ -72,6 +72,7 @@
     import anthropic
     import httpx
     import openai
+    import requests
     from django.http import HttpRequest, HttpResponse
     from fastapi import FastAPI
     from flask.app import Flask
@@ -84,7 +85,6 @@
     from .integrations.asgi import ASGIApp, ASGIInstrumentKwargs
     from .integrations.asyncpg import AsyncPGInstrumentKwargs
     from .integrations.aws_lambda import AwsLambdaInstrumentKwargs, LambdaHandler
-    from .integrations.celery import CeleryInstrumentKwargs
     from .integrations.flask import FlaskInstrumentKwargs
     from .integrations.httpx import AsyncClientKwargs, ClientKwargs, HTTPXInstrumentKwargs
     from .integrations.mysql import MySQLConnection, MySQLInstrumentKwargs
@@ -1224,17 +1224,26 @@ def instrument_httpx(
             **kwargs,
         )
 
-    def instrument_celery(self, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None:
+    def instrument_celery(self, **kwargs: Any) -> None:
         """Instrument `celery` so that spans are automatically created for each task.
 
         Uses the
         [OpenTelemetry Celery Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/celery/celery.html)
         library.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` method, for future compatibility.
         """
         from .integrations.celery import instrument_celery
 
         self._warn_if_not_initialized_for_instrumentation()
-        return instrument_celery(self, **kwargs)
+        return instrument_celery(
+            **{
+                'tracer_provider': self._config.get_tracer_provider(),
+                'meter_provider': self._config.get_meter_provider(),
+                **kwargs,
+            },
+        )
 
     def instrument_django(
         self,
@@ -1288,18 +1297,34 @@ def instrument_django(
             **kwargs,
         )
 
-    def instrument_requests(self, excluded_urls: str | None = None, **kwargs: Any) -> None:
+    def instrument_requests(
+        self,
+        excluded_urls: str | None = None,
+        request_hook: Callable[[Span, requests.PreparedRequest], None] | None = None,
+        response_hook: Callable[[Span, requests.PreparedRequest, requests.Response], None] | None = None,
+        **kwargs: Any,
+    ) -> None:
         """Instrument the `requests` module so that spans are automatically created for each request.
 
         Args:
             excluded_urls: A string containing a comma-delimited list of regexes used to exclude URLs from tracking
-            **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods,
-                particularly `request_hook` and `response_hook`.
+            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 `instrument` methods, for future compatibility.
         """
         from .integrations.requests import instrument_requests
 
         self._warn_if_not_initialized_for_instrumentation()
-        return instrument_requests(self, excluded_urls=excluded_urls, **kwargs)
+        return instrument_requests(
+            excluded_urls=excluded_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_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.