diff --git a/logfire-api/logfire_api/__init__.py b/logfire-api/logfire_api/__init__.py index 855c35c6d..7931b4374 100644 --- a/logfire-api/logfire_api/__init__.py +++ b/logfire-api/logfire_api/__init__.py @@ -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() diff --git a/logfire/_internal/instrument.py b/logfire/_internal/instrument.py index 3f80ceea0..28e961420 100644 --- a/logfire/_internal/instrument.py +++ b/logfire/_internal/instrument.py @@ -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 @@ -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 @@ -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: @@ -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. @@ -156,6 +172,8 @@ 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: + return NoopSpan() return get_logfire()._fast_span(final_span_name, attributes, **extra_span_kwargs()) # pyright: ignore[reportPrivateUsage] if extract_args is True: @@ -163,6 +181,8 @@ def open_span(*_: P.args, **__: P.kwargs): # pyright: ignore[reportRedeclaratio 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 @@ -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 diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index b2bb307e0..5b135c994 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -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. @@ -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 @@ -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. @@ -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( @@ -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 diff --git a/tests/test_logfire.py b/tests/test_logfire.py index a3a145558..ef79d6e28 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -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..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() == [] + + +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..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..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):