param version

Author: oliviaBahr

docs/guides/onboarding-checklist/add-manual-tracing.md

@@ -234,6 +234,17 @@ my_function(3, 4)
 # Logs: Applying my_function to x=3 and y=4
 ```
 
+You can also access the span directly within your instrumented function by adding a `logfire_span` parameter:
+
+```python
+@logfire.instrument('Processing {x=}')
+def process_data(x: int, logfire_span: logfire.LogfireSpan | None = None) -> int:
+    # Access and modify the span directly
+    if logfire_span:
+        logfire_span.message = f'Custom message for x={x}'
+    return x * 2
+```
+
 !!! note
 
     - The [`@logfire.instrument`][logfire.Logfire.instrument] decorator MUST be applied first, i.e., UNDER any other decorators.

logfire/_internal/instrument.py

@@ -63,19 +63,27 @@ def decorator(func: Callable[P, R]) -> Callable[P, R]:
         attributes = get_attributes(func, msg_template, tags)
         open_span = get_open_span(logfire, attributes, span_name, extract_args, func)
 
+        # Check if function has logfire_span parameter
+        sig = inspect.signature(func)
+        has_logfire_span_param = 'logfire_span' in sig.parameters
+
         if inspect.isgeneratorfunction(func):
             if not allow_generator:
                 warnings.warn(GENERATOR_WARNING_MESSAGE, stacklevel=2)
 
             def wrapper(*func_args: P.args, **func_kwargs: P.kwargs):  # type: ignore
-                with open_span(*func_args, **func_kwargs):
+                with open_span(*func_args, **func_kwargs) as span:
+                    if has_logfire_span_param:
+                        func_kwargs['logfire_span'] = span
                     yield from func(*func_args, **func_kwargs)
         elif inspect.isasyncgenfunction(func):
             if not allow_generator:
                 warnings.warn(GENERATOR_WARNING_MESSAGE, stacklevel=2)
 
             async def wrapper(*func_args: P.args, **func_kwargs: P.kwargs):  # type: ignore
-                with open_span(*func_args, **func_kwargs):
+                with open_span(*func_args, **func_kwargs) as span:
+                    if has_logfire_span_param:
+                        func_kwargs['logfire_span'] = span
                     # `yield from` is invalid syntax in an async function.
                     # This loop is not quite equivalent, because `yield from` also handles things like
                     # sending values to the subgenerator.
@@ -90,6 +98,8 @@ async def wrapper(*func_args: P.args, **func_kwargs: P.kwargs):  # type: ignore
 
             async def wrapper(*func_args: P.args, **func_kwargs: P.kwargs) -> R:  # type: ignore
                 with open_span(*func_args, **func_kwargs) as span:
+                    if has_logfire_span_param:
+                        func_kwargs['logfire_span'] = span
                     result = await func(*func_args, **func_kwargs)
                     if record_return:
                         # open_span returns a FastLogfireSpan, so we can't use span.set_attribute for complex types.
@@ -102,6 +112,8 @@ async def wrapper(*func_args: P.args, **func_kwargs: P.kwargs) -> R:  # type: ig
             # Same as the above, but without the async/await
             def wrapper(*func_args: P.args, **func_kwargs: P.kwargs) -> R:
                 with open_span(*func_args, **func_kwargs) as span:
+                    if has_logfire_span_param:
+                        func_kwargs['logfire_span'] = span
                     result = func(*func_args, **func_kwargs)
                     if record_return:
                         set_user_attributes_on_raw_span(span._span, {'return': result})
@@ -122,27 +134,32 @@ def get_open_span(
 ) -> Callable[P, AbstractContextManager[Any]]:
     final_span_name: str = span_name or attributes[ATTRIBUTES_MESSAGE_TEMPLATE_KEY]  # type: ignore
 
+    # Check if function has logfire_span parameter
+    sig = inspect.signature(func)
+    has_logfire_span_param = 'logfire_span' in sig.parameters
+
     # This is the fast case for when there are no arguments to extract
     def open_span(*_: P.args, **__: P.kwargs):  # type: ignore
+        if has_logfire_span_param:
+            return logfire._span(final_span_name, attributes)  # type: ignore
         return logfire._fast_span(final_span_name, attributes)  # type: ignore
 
     if extract_args is True:
-        sig = inspect.signature(func)
         if sig.parameters:  # only extract args if there are any
 
             def open_span(*func_args: P.args, **func_kwargs: P.kwargs):
                 bound = sig.bind(*func_args, **func_kwargs)
                 bound.apply_defaults()
                 args_dict = bound.arguments
+                if has_logfire_span_param:
+                    return logfire._span(final_span_name, {**attributes, **args_dict})  # type: ignore
                 return logfire._instrument_span_with_args(  # type: ignore
                     final_span_name, attributes, args_dict
                 )
 
         return open_span
 
     if extract_args:  # i.e. extract_args should be an iterable of argument names
-        sig = inspect.signature(func)
-
         if isinstance(extract_args, str):
             extract_args = [extract_args]
 
@@ -165,6 +182,8 @@ def open_span(*func_args: P.args, **func_kwargs: P.kwargs):
                 # This line is the only difference from the extract_args=True case
                 args_dict = {k: args_dict[k] for k in extract_args_final}
 
+                if has_logfire_span_param:
+                    return logfire._span(final_span_name, {**attributes, **args_dict})  # type: ignore
                 return logfire._instrument_span_with_args(  # type: ignore
                     final_span_name, attributes, args_dict
                 )

tests/test_logfire.py

@@ -1273,6 +1273,66 @@ def run(a: str) -> Model:
     )
 
 
+def test_instrument_with_logfire_span_parameter(exporter: TestExporter):
+    @logfire.instrument('Calling foo with {x=}')
+    def foo(x: int, logfire_span: logfire.LogfireSpan | None = None) -> int:
+        # Test that we can access the span and modify its message
+        assert logfire_span is not None
+        logfire_span.message = f'Modified message for x={x}'
+        return x * 2
+
+    result = foo(5)
+    assert result == 10
+
+    spans = exporter.exported_spans_as_dict(_strip_function_qualname=False)
+    assert len(spans) == 1
+    span = spans[0]
+    assert span['attributes']['logfire.msg'] == 'Modified message for x=5'
+    assert span['attributes']['x'] == 5
+
+
+def test_instrument_with_logfire_span_parameter_async(exporter: TestExporter):
+    @logfire.instrument('Calling async foo with {x=}')
+    async def foo(x: int, logfire_span: logfire.LogfireSpan | None = None) -> int:
+        # Test that we can access the span and modify its message
+        assert logfire_span is not None
+        logfire_span.message = f'Async modified message for x={x}'
+        return x * 3
+
+    async def run_test():
+        return await foo(7)
+
+    import asyncio
+
+    result = asyncio.run(run_test())
+    assert result == 21
+
+    spans = exporter.exported_spans_as_dict(_strip_function_qualname=False)
+    assert len(spans) == 1
+    span = spans[0]
+    assert span['attributes']['logfire.msg'] == 'Async modified message for x=7'
+    assert span['attributes']['x'] == 7
+
+
+def test_instrument_with_logfire_span_parameter_extract_args_false(exporter: TestExporter):
+    @logfire.instrument('Calling foo', extract_args=False)
+    def foo(x: int, logfire_span: logfire.LogfireSpan | None = None) -> int:
+        # Test that we can access the span and modify its message
+        assert logfire_span is not None
+        logfire_span.message = f'Extract args false message for x={x}'
+        return x * 4
+
+    result = foo(3)
+    assert result == 12
+
+    spans = exporter.exported_spans_as_dict(_strip_function_qualname=False)
+    assert len(spans) == 1
+    span = spans[0]
+    assert span['attributes']['logfire.msg'] == 'Extract args false message for x=3'
+    # x should not be in attributes since extract_args=False
+    assert 'x' not in span['attributes']
+
+
 def test_validation_error_on_span(exporter: TestExporter) -> None:
     class Model(BaseModel, plugin_settings={'logfire': {'record': 'off'}}):
         a: int