Skip to content

Add new_context to instrumentation #1499

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Add new_context to instrumentation #1499

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
33ac479
Add new_context to instrumentation
njz-cvm Oct 24, 2025
315dc95
Add _strip_function_qualname to testing
njz-cvm Oct 24, 2025
d43e4b8
Removed factory option, added linking, fixed logic error
njz-cvm Oct 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 15 additions & 4 deletions logfire/_internal/instrument.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
from .utils import safe_repr, uniquify_sequence

if TYPE_CHECKING:
from logfire.propagate import ContextCarrier

from .main import Logfire


Expand Down Expand Up @@ -50,6 +52,7 @@ def instrument(
extract_args: bool | Iterable[str],
record_return: bool,
allow_generator: bool,
new_context: bool | Callable[[], ContextCarrier],
) -> Callable[[Callable[P, R]], Callable[P, R]]:
from .main import set_user_attributes_on_raw_span

Expand All @@ -61,7 +64,7 @@ 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)
open_span = get_open_span(logfire, attributes, span_name, extract_args, func, new_context)

if inspect.isgeneratorfunction(func):
if not allow_generator:
Expand Down Expand Up @@ -119,7 +122,10 @@ def get_open_span(
span_name: str | None,
extract_args: bool | Iterable[str],
func: Callable[P, R],
new_context: bool | Callable[[], ContextCarrier],
) -> Callable[P, AbstractContextManager[Any]]:
from logfire.propagate import attach_context

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
Expand All @@ -138,9 +144,7 @@ def open_span(*func_args: P.args, **func_kwargs: P.kwargs):
final_span_name, attributes, args_dict
)

return open_span

if extract_args: # i.e. extract_args should be an iterable of argument names
elif extract_args: # i.e. extract_args should be an iterable of argument names
sig = inspect.signature(func)

if isinstance(extract_args, str):
Expand Down Expand Up @@ -169,6 +173,13 @@ def open_span(*func_args: P.args, **func_kwargs: P.kwargs):
final_span_name, attributes, args_dict
)

if new_context is not False:
_open_span = open_span

def open_span(*func_args: P.args, **func_kwargs: P.kwargs):
with attach_context({} if new_context is True else new_context()):
return _open_span(*func_args, **func_kwargs)

return open_span


Expand Down
10 changes: 9 additions & 1 deletion logfire/_internal/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,8 @@
from starlette.websockets import WebSocket
from typing_extensions import Unpack

from logfire.propagate import ContextCarrier

from ..integrations.aiohttp_client import (
RequestHook as AiohttpClientRequestHook,
ResponseHook as AiohttpClientResponseHook,
Expand Down Expand Up @@ -580,6 +582,7 @@ def instrument(
extract_args: bool | Iterable[str] = True,
record_return: bool = False,
allow_generator: bool = False,
new_context: bool | Callable[[], ContextCarrier] = False,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
"""Decorator for instrumenting a function as a span.

Expand All @@ -603,6 +606,8 @@ def my_function(a: int):
Ignored for generators.
allow_generator: Set to `True` to prevent a warning when instrumenting a generator function.
Read https://logfire.pydantic.dev/docs/guides/advanced/generators/#using-logfireinstrument first.
new_context: Set to `True` to clear context before starting instrumentation.
A callable can instead be used to provide starting context. `new_context=True` is the same as `new_context=dict`.
"""

@overload
Expand All @@ -629,6 +634,7 @@ def instrument( # type: ignore[reportInconsistentOverload]
extract_args: bool | Iterable[str] = True,
record_return: bool = False,
allow_generator: bool = False,
new_context: bool | Callable[[], ContextCarrier] = False,
) -> Callable[[Callable[P, R]], Callable[P, R]] | Callable[P, R]:
"""Decorator for instrumenting a function as a span.

Expand All @@ -652,11 +658,13 @@ def my_function(a: int):
Ignored for generators.
allow_generator: Set to `True` to prevent a warning when instrumenting a generator function.
Read https://logfire.pydantic.dev/docs/guides/advanced/generators/#using-logfireinstrument first.
new_context: Set to `True` to clear context before starting instrumentation.
A callable can instead be used to provide starting context. `new_context=True` is the same as `new_context=dict`.
"""
if callable(msg_template):
return self.instrument()(msg_template)
return instrument(
self, tuple(self._tags), msg_template, span_name, extract_args, record_return, allow_generator
self, tuple(self._tags), msg_template, span_name, extract_args, record_return, allow_generator, new_context
)

def log(
Expand Down
207 changes: 207 additions & 0 deletions tests/test_logfire.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -748,6 +748,213 @@ def hello_world(a: int) -> str:
)


def test_instrument_with_parent(exporter: TestExporter) -> None:
tagged = logfire.with_tags('test_instrument')

@tagged.instrument('hello-world {a=}', record_return=True)
def hello_world(a: int) -> str:
return f'hello {a}'

@tagged.instrument('parent', record_return=True)
def parent() -> str:
return hello_world(5)

assert parent() == 'hello 5'

assert exporter.exported_spans_as_dict(_include_pending_spans=True, _strip_function_qualname=False) == snapshot(
[
{
'name': 'parent',
'context': {'trace_id': 1, 'span_id': 2, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 1000000000,
'end_time': 1000000000,
'attributes': {
'code.filepath': 'test_logfire.py',
'code.lineno': 123,
'code.function': 'test_instrument_with_parent.<locals>.parent',
'logfire.msg': 'parent',
'logfire.msg_template': 'parent',
'logfire.span_type': 'pending_span',
'logfire.pending_parent_id': '0000000000000000',
'logfire.tags': ('test_instrument',),
},
},
{
'name': 'hello-world {a=}',
'context': {'trace_id': 1, 'span_id': 4, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
'start_time': 2000000000,
'end_time': 2000000000,
'attributes': {
'code.filepath': 'test_logfire.py',
'code.lineno': 123,
'code.function': 'test_instrument_with_parent.<locals>.hello_world',
'a': 5,
'logfire.msg_template': 'hello-world {a=}',
'logfire.msg': 'hello-world a=5',
'logfire.json_schema': '{"type":"object","properties":{"a":{}}}',
'logfire.span_type': 'pending_span',
'logfire.pending_parent_id': '0000000000000001',
'logfire.tags': ('test_instrument',),
},
},
{
'attributes': {
'a': 5,
'code.filepath': 'test_logfire.py',
'code.function': 'test_instrument_with_parent.<locals>.hello_world',
'code.lineno': 123,
'logfire.json_schema': '{"type":"object","properties":{"a":{},"return":{}}}',
'logfire.msg': 'hello-world a=5',
'logfire.msg_template': 'hello-world {a=}',
'logfire.span_type': 'span',
'logfire.tags': ('test_instrument',),
'return': 'hello 5',
},
'context': {
'is_remote': False,
'span_id': 3,
'trace_id': 1,
},
'end_time': 3000000000,
'name': 'hello-world {a=}',
'parent': {
'is_remote': False,
'span_id': 1,
'trace_id': 1,
},
'start_time': 2000000000,
},
{
'attributes': {
'code.filepath': 'test_logfire.py',
'code.function': 'test_instrument_with_parent.<locals>.parent',
'code.lineno': 123,
'logfire.json_schema': '{"type":"object","properties":{"return":{}}}',
'logfire.msg': 'parent',
'logfire.msg_template': 'parent',
'logfire.span_type': 'span',
'logfire.tags': ('test_instrument',),
'return': 'hello 5',
},
'context': {
'is_remote': False,
'span_id': 1,
'trace_id': 1,
},
'end_time': 4000000000,
'name': 'parent',
'parent': None,
'start_time': 1000000000,
},
]
)


def test_instrument_new_context(exporter: TestExporter) -> None:
tagged = logfire.with_tags('test_instrument')

def context_factory():
return {'new_contex': 123}
Copy link
Contributor

@alexmojaki alexmojaki Oct 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We didn't discuss this context factory API, I'm not sure if I like it or understand the point. It's certainly not being tested here, new_contex doesn't do anything or appear in the exported spans.

Copy link
Author

@njz-cvm njz-cvm Oct 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't mention it in the original issue because it wasn't something I'd considered. But when I realized I'd just be using attach_context behind the scenes, it seemed preferable to allow users the option to provide input for attach_context.

If you're against it, it's not that useful for me and I don't mind cutting it. But if it's worth keeping I'll add better testing for it.


@tagged.instrument('hello-world {a=}', record_return=True, new_context=context_factory)
def hello_world(a: int) -> str:
return f'hello {a}'

@tagged.instrument('parent', record_return=True)
def parent() -> str:
return hello_world(5)

assert parent() == 'hello 5'

assert exporter.exported_spans_as_dict(_include_pending_spans=True, _strip_function_qualname=False) == snapshot(
[
{
'name': 'parent',
'context': {'trace_id': 1, 'span_id': 2, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 1000000000,
'end_time': 1000000000,
'attributes': {
'code.filepath': 'test_logfire.py',
'code.lineno': 123,
'code.function': 'test_instrument_new_context.<locals>.parent',
'logfire.msg': 'parent',
'logfire.msg_template': 'parent',
'logfire.span_type': 'pending_span',
'logfire.pending_parent_id': '0000000000000000',
'logfire.tags': ('test_instrument',),
},
},
{
'name': 'hello-world {a=}',
'context': {'trace_id': 2, 'span_id': 4, 'is_remote': False},
'parent': {'trace_id': 2, 'span_id': 3, 'is_remote': False},
'start_time': 2000000000,
'end_time': 2000000000,
'attributes': {
'code.filepath': 'test_logfire.py',
'code.lineno': 123,
'code.function': 'test_instrument_new_context.<locals>.hello_world',
'a': 5,
'logfire.msg_template': 'hello-world {a=}',
'logfire.msg': 'hello-world a=5',
'logfire.json_schema': '{"type":"object","properties":{"a":{}}}',
'logfire.span_type': 'pending_span',
'logfire.pending_parent_id': '0000000000000000',
'logfire.tags': ('test_instrument',),
},
},
{
'attributes': {
'a': 5,
'code.filepath': 'test_logfire.py',
'code.function': 'test_instrument_new_context.<locals>.hello_world',
'code.lineno': 123,
'logfire.json_schema': '{"type":"object","properties":{"a":{},"return":{}}}',
'logfire.msg': 'hello-world a=5',
'logfire.msg_template': 'hello-world {a=}',
'logfire.span_type': 'span',
'logfire.tags': ('test_instrument',),
'return': 'hello 5',
},
'context': {
'is_remote': False,
'span_id': 3,
'trace_id': 2,
},
'end_time': 3000000000,
'name': 'hello-world {a=}',
'parent': None,
'start_time': 2000000000,
},
{
'attributes': {
'code.filepath': 'test_logfire.py',
'code.function': 'test_instrument_new_context.<locals>.parent',
'code.lineno': 123,
'logfire.json_schema': '{"type":"object","properties":{"return":{}}}',
'logfire.msg': 'parent',
'logfire.msg_template': 'parent',
'logfire.span_type': 'span',
'logfire.tags': ('test_instrument',),
'return': 'hello 5',
},
'context': {
'is_remote': False,
'span_id': 1,
'trace_id': 1,
},
'end_time': 4000000000,
'name': 'parent',
'parent': None,
'start_time': 1000000000,
},
]
)


def test_instrument_other_callable(exporter: TestExporter):
class Instrumented:
def __call__(self, a: int) -> str:
Expand Down
Loading