Skip to content
Merged
Show file tree
Hide file tree
Changes from 6 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 @@ -131,7 +131,7 @@ def log_slow_async_callbacks(self, *args, **kwargs) -> None: # pragma: no cover

def install_auto_tracing(self, *args, **kwargs) -> None: ...

def instrument(self, *args, **kwargs):
def instrument(self, *args, _level=None, **kwargs):
Comment thread
imp-joshi marked this conversation as resolved.
Outdated
def decorator(func):
return func

Comment thread
alexmojaki marked this conversation as resolved.
Outdated
Expand Down
30 changes: 28 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,16 @@ 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]]:
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}
Comment thread
alexmojaki marked this conversation as resolved.
Outdated

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 +170,21 @@ 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.
from .main import NoopSpan

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:
from .main import NoopSpan

return NoopSpan()
bound = sig.bind(*func_args, **func_kwargs)
bound.apply_defaults()
args_dict = bound.arguments
Expand Down Expand Up @@ -190,6 +212,10 @@ 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:
from .main import NoopSpan
Comment thread
imp-joshi marked this conversation as resolved.
Outdated

return NoopSpan()
bound = sig.bind(*func_args, **func_kwargs)
bound.apply_defaults()
args_dict = bound.arguments
Expand Down
21 changes: 20 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,
Comment thread
imp-joshi marked this conversation as resolved.
Outdated
) -> 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,11 @@ def message(self) -> str: # pragma: no cover
def message(self, message: str):
pass

@property
def _span(self) -> Any:
# set_user_attributes_on_raw_span checks is_recording() first; INVALID_SPAN returns False
Comment thread
imp-joshi marked this conversation as resolved.
Outdated
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