Skip to content
Merged
Show file tree
Hide file tree
Changes from 9 commits
Commits
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
2 changes: 1 addition & 1 deletion logfire-api/logfire_api/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -322,5 +322,5 @@ def set_baggage(*args, **kwargs) -> ContextManager[None]:
def get_context(*args, **kwargs) -> dict[str, Any]:
return {}

def attach_context(*args, **kwargs)-> ContextManager[None]:
def attach_context(*args, **kwargs) -> ContextManager[None]:
return nullcontext()
26 changes: 24 additions & 2 deletions logfire/_internal/instrument.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,13 @@
from opentelemetry.util import types as otel_types
from typing_extensions import LiteralString, ParamSpec

from .constants import ATTRIBUTES_MESSAGE_TEMPLATE_KEY, ATTRIBUTES_TAGS_KEY
from .constants import (
ATTRIBUTES_LOG_LEVEL_NUM_KEY,
ATTRIBUTES_MESSAGE_TEMPLATE_KEY,
ATTRIBUTES_TAGS_KEY,
LevelName,
log_level_attributes,
)
from .stack_info import get_filepath_attribute
from .utils import safe_repr, uniquify_sequence

Expand Down Expand Up @@ -52,6 +58,7 @@ def instrument(
record_return: bool,
allow_generator: bool,
new_trace: bool,
level: LevelName | int | None = None,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
from .main import set_user_attributes_on_raw_span

Expand All @@ -63,7 +70,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, new_trace)
open_span = get_open_span(logfire, attributes, span_name, extract_args, func, new_trace, level=level)

if inspect.isgeneratorfunction(func):
if not allow_generator:
Expand Down Expand Up @@ -122,9 +129,18 @@ def get_open_span(
extract_args: bool | Iterable[str],
func: Callable[P, R],
new_trace: bool,
level: LevelName | int | None = None,
) -> Callable[P, AbstractContextManager[Any]]:
from .main import NoopSpan

final_span_name: str = span_name or attributes[ATTRIBUTES_MESSAGE_TEMPLATE_KEY] # pyright: ignore[reportAssignmentType]

level_num: int | None = None
if level is not None:
level_attrs = log_level_attributes(level)
level_num = int(level_attrs[ATTRIBUTES_LOG_LEVEL_NUM_KEY])
attributes = {**attributes, **level_attrs}

def get_logfire():
# This avoids having a `logfire` closure variable, which would make the instrumented
# function unpicklable with cloudpickle.
Expand Down Expand Up @@ -156,13 +172,17 @@ def extra_span_kwargs() -> dict[str, Any]:

# This is the fast case for when there are no arguments to extract
def open_span(*_: P.args, **__: P.kwargs): # pyright: ignore[reportRedeclaration]
if level_num is not None and level_num < get_logfire().config.min_level:
Comment thread
imp-joshi marked this conversation as resolved.
return NoopSpan()
return get_logfire()._fast_span(final_span_name, attributes, **extra_span_kwargs()) # pyright: ignore[reportPrivateUsage]

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):
if level_num is not None and level_num < get_logfire().config.min_level:
return NoopSpan()
bound = sig.bind(*func_args, **func_kwargs)
bound.apply_defaults()
args_dict = bound.arguments
Expand Down Expand Up @@ -190,6 +210,8 @@ def open_span(*func_args: P.args, **func_kwargs: P.kwargs):
if extract_args_final: # check that there are still arguments to extract

def open_span(*func_args: P.args, **func_kwargs: P.kwargs):
if level_num is not None and level_num < get_logfire().config.min_level:
return NoopSpan()
bound = sig.bind(*func_args, **func_kwargs)
bound.apply_defaults()
args_dict = bound.arguments
Expand Down
20 changes: 19 additions & 1 deletion logfire/_internal/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -608,6 +608,7 @@ def instrument(
record_return: bool = False,
allow_generator: bool = False,
new_trace: bool = False,
level: LevelName | int | None = None,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
"""Decorator for instrumenting a function as a span.

Expand All @@ -633,6 +634,8 @@ def my_function(a: int):
Read https://logfire.pydantic.dev/docs/guides/advanced/generators/#using-logfireinstrument first.
new_trace: Set to `True` to start a new trace with a span link to the current span
instead of creating a child of the current span.
level: The log level for the span. If provided, the span will be tagged with this level
and suppressed if the level is below the configured `min_level`.
"""

@overload
Expand Down Expand Up @@ -660,6 +663,7 @@ def instrument( # pyright: ignore[reportInconsistentOverload]
record_return: bool = False,
allow_generator: bool = False,
new_trace: bool = False,
level: LevelName | int | None = None,
) -> Callable[[Callable[P, R]], Callable[P, R]] | Callable[P, R]:
"""Decorator for instrumenting a function as a span.

Expand All @@ -685,11 +689,21 @@ def my_function(a: int):
Read https://logfire.pydantic.dev/docs/guides/advanced/generators/#using-logfireinstrument first.
new_trace: Set to `True` to start a new trace with a span link to the current span
instead of creating a child of the current span.
level: The log level for the span. If provided, the span will be tagged with this level
and suppressed if the level is below the configured `min_level`.
"""
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, new_trace
self,
tuple(self._tags),
msg_template,
span_name,
extract_args,
record_return,
allow_generator,
new_trace,
level=level,
)

def log(
Expand Down Expand Up @@ -3083,6 +3097,10 @@ def message(self) -> str: # pragma: no cover
def message(self, message: str):
pass

@property
def _span(self) -> Any:
return trace_api.INVALID_SPAN

def is_recording(self) -> bool:
return False

Expand Down
135 changes: 135 additions & 0 deletions tests/test_logfire.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,141 @@ def test_log_methods_without_kwargs(method: str):
""")


def test_instrument_with_level(exporter: TestExporter) -> None:
@logfire.instrument('my span', level='warn', extract_args=False)
def my_func() -> str:
return 'ok'

assert my_func() == 'ok'
assert exporter.exported_spans_as_dict(_strip_function_qualname=False) == snapshot(
[
{
'name': 'my span',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': 1000000000,
'end_time': 2000000000,
'attributes': {
'code.function': 'test_instrument_with_level.<locals>.my_func',
'logfire.msg_template': 'my span',
'code.lineno': 123,
'code.filepath': 'test_logfire.py',
'logfire.level_num': 13,
'logfire.span_type': 'span',
'logfire.msg': 'my span',
},
}
]
)


def test_instrument_level_filtered(exporter: TestExporter, config_kwargs: dict[str, Any]) -> None:
config_kwargs['min_level'] = 'info'
logfire.configure(**config_kwargs)

@logfire.instrument('my span', level='debug', extract_args=False)
def my_func() -> str:
return 'ok'

assert my_func() == 'ok'
assert exporter.exported_spans_as_dict() == []


def test_instrument_level_filtered_record_return(exporter: TestExporter, config_kwargs: dict[str, Any]) -> None:
config_kwargs['min_level'] = 'info'
logfire.configure(**config_kwargs)

@logfire.instrument('my span', level='debug', extract_args=False, record_return=True)
def my_func() -> str:
return 'ok'

assert my_func() == 'ok'
assert exporter.exported_spans_as_dict() == []
Comment thread
imp-joshi marked this conversation as resolved.


def test_instrument_level_filtered_extract_args(exporter: TestExporter, config_kwargs: dict[str, Any]) -> None:
config_kwargs['min_level'] = 'info'
logfire.configure(**config_kwargs)

@logfire.instrument('my span {x=}', level='debug', extract_args=True)
def my_func(x: int) -> int:
return x * 2

assert my_func(5) == 10
assert exporter.exported_spans_as_dict() == []


def test_instrument_level_filtered_extract_args_iterable(exporter: TestExporter, config_kwargs: dict[str, Any]) -> None:
config_kwargs['min_level'] = 'info'
logfire.configure(**config_kwargs)

@logfire.instrument('my span', level='debug', extract_args=('x',))
def my_func(x: int) -> int:
return x * 2

assert my_func(5) == 10
assert exporter.exported_spans_as_dict() == []


@pytest.mark.anyio
async def test_instrument_with_level_async(exporter: TestExporter) -> None:
@logfire.instrument('async span', level='warn', extract_args=False)
async def my_async_func() -> str:
return 'ok'

assert await my_async_func() == 'ok'
assert exporter.exported_spans_as_dict(_strip_function_qualname=False) == snapshot(
[
{
'name': 'async span',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': 1000000000,
'end_time': 2000000000,
'attributes': {
'code.function': 'test_instrument_with_level_async.<locals>.my_async_func',
'logfire.msg_template': 'async span',
'code.lineno': 123,
'code.filepath': 'test_logfire.py',
'logfire.level_num': 13,
'logfire.span_type': 'span',
'logfire.msg': 'async span',
},
}
]
)


def test_instrument_with_level_and_extract_args(exporter: TestExporter) -> None:
@logfire.instrument('span {x=}', level='warn')
def my_func(x: int) -> int:
return x * 2

assert my_func(5) == 10
assert exporter.exported_spans_as_dict(_strip_function_qualname=False) == snapshot(
[
{
'name': 'span {x=}',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': 1000000000,
'end_time': 2000000000,
'attributes': {
'code.function': 'test_instrument_with_level_and_extract_args.<locals>.my_func',
'logfire.msg_template': 'span {x=}',
'code.lineno': 123,
'code.filepath': 'test_logfire.py',
'logfire.level_num': 13,
'logfire.msg': 'span x=5',
'logfire.json_schema': '{"type":"object","properties":{"x":{}}}',
'x': 5,
'logfire.span_type': 'span',
},
}
]
)


def test_instrument_with_no_args(exporter: TestExporter) -> None:
@logfire.instrument()
def foo(x: int):
Expand Down
Loading