context var version

Author: oliviaBahr

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

@@ -234,6 +234,15 @@ my_function(3, 4)
 # Logs: Applying my_function to x=3 and y=4
 ```
 
+Access the created span inside the function using `logfire.current_span()`:
+
+```python
+@logfire.instrument('Processing')
+def process_data(user_dict: dict):
+    user_id = user_dict.get("id")
+    logfire.current_span().message = f'Processing User: {user_id}'
+```
+
 !!! note
 
     - The [`@logfire.instrument`][logfire.Logfire.instrument] decorator MUST be applied first, i.e., UNDER any other decorators.

logfire/__init__.py

@@ -24,6 +24,7 @@
 DEFAULT_LOGFIRE_INSTANCE: Logfire = Logfire()
 span = DEFAULT_LOGFIRE_INSTANCE.span
 instrument = DEFAULT_LOGFIRE_INSTANCE.instrument
+current_span = DEFAULT_LOGFIRE_INSTANCE.current_span
 force_flush = DEFAULT_LOGFIRE_INSTANCE.force_flush
 log_slow_async_callbacks = DEFAULT_LOGFIRE_INSTANCE.log_slow_async_callbacks
 install_auto_tracing = DEFAULT_LOGFIRE_INSTANCE.install_auto_tracing
@@ -108,6 +109,7 @@ def loguru_handler() -> Any:
     'configure',
     'span',
     'instrument',
+    'current_span',
     'log',
     'trace',
     'debug',

logfire/_internal/ast_utils.py

@@ -4,6 +4,7 @@
 import functools
 import inspect
 import sys
+import textwrap
 import types
 import warnings
 from abc import ABC, abstractmethod
@@ -245,6 +246,25 @@ class InspectArgumentsFailedWarning(Warning):
     pass
 
 
+@functools.lru_cache(maxsize=1024)
+def has_current_span_call(func: Any) -> bool:
+    """Check if a function contains calls to logfire.current_span()."""
+    try:
+        tree = ast.parse(textwrap.dedent(inspect.getsource(func)))
+        for node in ast.walk(tree):
+            if (
+                isinstance(node, ast.Call)
+                and isinstance(node.func, ast.Attribute)
+                and isinstance(node.func.value, ast.Name)
+                and node.func.attr == 'current_span'
+                and node.func.value.id == 'logfire'
+            ):
+                return True
+        return False
+    except (OSError, SyntaxError, TypeError, AttributeError):
+        return False
+
+
 @functools.lru_cache(maxsize=1024)
 def get_node_source_text(node: ast.AST, ex_source: executing.Source):
     """Returns some Python source code representing `node`.

logfire/_internal/instrument.py

@@ -11,6 +11,7 @@
 from opentelemetry.util import types as otel_types
 from typing_extensions import LiteralString, ParamSpec
 
+from .ast_utils import has_current_span_call
 from .constants import ATTRIBUTES_MESSAGE_TEMPLATE_KEY, ATTRIBUTES_TAGS_KEY
 from .stack_info import get_filepath_attribute
 from .utils import safe_repr, uniquify_sequence
@@ -61,7 +62,8 @@ 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)
+        uses_current_span = has_current_span_call(func)
+        open_span = get_open_span(logfire, attributes, span_name, extract_args, uses_current_span, func)
 
         if inspect.isgeneratorfunction(func):
             if not allow_generator:
@@ -90,21 +92,31 @@ 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:
+                    token = None
+                    if uses_current_span:
+                        token = logfire._current_span_var.set(span)  # type: ignore[protected-access]
                     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.
                         # This isn't great because it has to parse the JSON schema.
                         # Not sure if making get_open_span return a LogfireSpan when record_return is True
                         # would be faster overall or if it would be worth the added complexity.
                         set_user_attributes_on_raw_span(span._span, {'return': result})
+                    if token:
+                        logfire._current_span_var.reset(token)  # type: ignore[protected-access]
                     return result
         else:
             # 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:
+                    token = None
+                    if uses_current_span:
+                        token = logfire._current_span_var.set(span)  # type: ignore[protected-access]
                     result = func(*func_args, **func_kwargs)
                     if record_return:
                         set_user_attributes_on_raw_span(span._span, {'return': result})
+                    if token:
+                        logfire._current_span_var.reset(token)  # type: ignore[protected-access]
                     return result
 
         wrapper = functools.wraps(func)(wrapper)  # type: ignore
@@ -118,12 +130,15 @@ def get_open_span(
     attributes: dict[str, otel_types.AttributeValue],
     span_name: str | None,
     extract_args: bool | Iterable[str],
+    uses_current_span: bool,
     func: Callable[P, R],
 ) -> Callable[P, AbstractContextManager[Any]]:
     final_span_name: str = span_name or attributes[ATTRIBUTES_MESSAGE_TEMPLATE_KEY]  # type: ignore
 
     # This is the fast case for when there are no arguments to extract
     def open_span(*_: P.args, **__: P.kwargs):  # type: ignore
+        if uses_current_span:
+            return logfire._span(final_span_name, attributes)  # type: ignore[protected-access]
         return logfire._fast_span(final_span_name, attributes)  # type: ignore
 
     if extract_args is True:
@@ -134,6 +149,9 @@ 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 uses_current_span:
+                    return logfire._span(final_span_name, {**attributes, **args_dict})  # type: ignore[protected-access]
+
                 return logfire._instrument_span_with_args(  # type: ignore
                     final_span_name, attributes, args_dict
                 )
@@ -165,6 +183,9 @@ 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 uses_current_span:
+                    return logfire._span(final_span_name, {**attributes, **args_dict})  # type: ignore[protected-access]
+
                 return logfire._instrument_span_with_args(  # type: ignore
                     final_span_name, attributes, args_dict
                 )

logfire/_internal/main.py

@@ -7,7 +7,7 @@
 import warnings
 from collections.abc import Iterable, Sequence
 from contextlib import AbstractContextManager
-from contextvars import Token
+from contextvars import ContextVar, Token
 from enum import Enum
 from functools import cached_property
 from time import time
@@ -138,6 +138,7 @@ def __init__(
         self._sample_rate = sample_rate
         self._console_log = console_log
         self._otel_scope = otel_scope
+        self._current_span_var = ContextVar('logfire_current_span', default=NoopSpan())
 
     @property
     def config(self) -> LogfireConfig:
@@ -171,6 +172,9 @@ def _get_tracer(self, *, is_span_tracer: bool) -> Tracer:  # pragma: no cover
             is_span_tracer=is_span_tracer,
         )
 
+    def current_span(self) -> LogfireSpan:
+        return self._current_span_var.get()  # type: ignore[return-value]
+
     # If any changes are made to this method, they may need to be reflected in `_fast_span` as well.
     def _span(
         self,

tests/test_logfire.py

@@ -38,7 +38,7 @@
     LevelName,
 )
 from logfire._internal.formatter import FormattingFailedWarning
-from logfire._internal.main import NoopSpan
+from logfire._internal.main import LogfireSpan, NoopSpan
 from logfire._internal.tracer import record_exception
 from logfire._internal.utils import SeededRandomIdGenerator, is_instrumentation_suppressed
 from logfire.integrations.logging import LogfireLoggingHandler
@@ -1344,6 +1344,83 @@ def run(a: str) -> None:
     )
 
 
+def test_current_span_default_is_noop():
+    span = logfire.DEFAULT_LOGFIRE_INSTANCE.current_span()
+    assert isinstance(span, NoopSpan)
+
+
+def test_current_span_nested(exporter: TestExporter):
+    @logfire.instrument('outer')
+    def outer():
+        s = logfire.current_span()
+        s.message = 'Starting outer operation'
+        inner()
+        assert logfire.current_span() == s
+        s = logfire.current_span()
+        s.message = 'Completing outer operation'
+
+    @logfire.instrument('inner')
+    def inner():
+        s = logfire.current_span()
+        s.message = 'Processing inner operation'
+
+    outer()
+    assert isinstance(logfire.current_span(), NoopSpan)
+
+    spans = exporter.exported_spans_as_dict(_strip_function_qualname=False)
+
+    assert len(spans) == 2
+    assert spans[0]['name'] == 'inner'
+    assert spans[0]['attributes']['logfire.msg'] == 'Processing inner operation'
+    assert spans[1]['name'] == 'outer'
+    assert spans[1]['attributes']['logfire.msg'] == 'Completing outer operation'
+
+
+@pytest.mark.anyio
+async def test_current_span_async(exporter: TestExporter):
+    print('Testing current_span_async - functions should auto-detect current_span usage')
+
+    @logfire.instrument('async outer')
+    async def outer():
+        s = logfire.current_span()
+        assert isinstance(s, LogfireSpan)
+        s.message = 'Starting async outer operation'
+        await inner()
+        s = logfire.current_span()
+        assert isinstance(s, LogfireSpan)
+        s.message = 'Completing async outer operation'
+
+    @logfire.instrument('async inner')
+    async def inner():
+        s = logfire.current_span()
+        assert isinstance(s, LogfireSpan)
+        s.message = 'Processing async inner operation'
+
+    await outer()
+    assert isinstance(logfire.current_span(), NoopSpan)
+
+    spans = exporter.exported_spans_as_dict(_strip_function_qualname=False)
+
+    assert len(spans) == 2
+    assert spans[0]['name'] == 'async inner'
+    assert spans[0]['attributes']['logfire.msg'] == 'Processing async inner operation'
+    assert spans[1]['name'] == 'async outer'
+    assert spans[1]['attributes']['logfire.msg'] == 'Completing async outer operation'
+
+
+def test_fast_span_when_current_span_not_called(exporter: TestExporter):
+    """Test that when current_span() is not called, the exported span is a fast span"""
+
+    @logfire.instrument
+    def fast_operation(): ...
+
+    fast_operation()
+    spans = exporter.exported_spans_as_dict(_strip_function_qualname=False)
+
+    # pending_span would be a LogfireSpan, just span is FastLogfireSpan
+    assert spans[0]['attributes']['logfire.span_type'] == 'span'
+
+
 @dataclass
 class Foo:
     x: int