Warn user if f-string expression contains await (#944)

Co-authored-by: Alex Hall <[email protected]>

Author: Sbargaoui

logfire/_internal/formatter.py

@@ -336,6 +336,8 @@ def logfire_format_with_magic(
         return ''.join(chunk['v'] for chunk in chunks), extra_attrs, new_template
     except KnownFormattingError as e:
         warn_formatting(str(e) or str(e.__cause__))
+    except FStringAwaitError as e:
+        warn_fstring_await(str(e))
     except Exception:
         # This is an unexpected error that likely indicates a bug in our logic.
         # Handle it here so that the span/log still gets created, just without a nice message.
@@ -354,6 +356,12 @@ def compile_formatted_value(node: ast.FormattedValue, ex_source: executing.Sourc
     3. Another code object which formats the value.
     """
     source = get_node_source_text(node.value, ex_source)
+
+    # Check if the expression contains await before attempting to compile
+    for sub_node in ast.walk(node.value):
+        if isinstance(sub_node, ast.Await):
+            raise FStringAwaitError(source)
+
     value_code = compile(source, '<fvalue1>', 'eval')
     expr = ast.Expression(
         ast.JoinedStr(
@@ -435,6 +443,14 @@ class KnownFormattingError(Exception):
     """
 
 
+class FStringAwaitError(Exception):
+    """An error raised when an await expression is found in an f-string.
+
+    This is a specific case that can't be handled by f-string introspection and requires
+    pre-evaluating the await expression before logging.
+    """
+
+
 class FormattingFailedWarning(UserWarning):
     pass
 
@@ -449,3 +465,17 @@ def warn_formatting(msg: str):
         f'    The problem was: {msg}',
         category=FormattingFailedWarning,
     )
+
+
+def warn_fstring_await(msg: str):
+    warn_at_user_stacklevel(
+        f'\n'
+        f'    Cannot evaluate await expression in f-string. Pre-evaluate the expression before logging.\n'
+        '    For example, change:\n'
+        '      logfire.info(f"{await get_value()}")\n'
+        '    To:\n'
+        '      value = await get_value()\n'
+        '      logfire.info(f"{value}")\n'
+        f'    The problematic f-string value was: {msg}',
+        category=FormattingFailedWarning,
+    )

tests/test_formatter.py

@@ -1,13 +1,16 @@
 import contextlib
+import sys
 from collections import ChainMap
 from types import SimpleNamespace
 from typing import Any, Mapping
 
 import pytest
 from inline_snapshot import snapshot
 
+import logfire
 from logfire._internal.formatter import FormattingFailedWarning, chunks_formatter, logfire_format
 from logfire._internal.scrubbing import NOOP_SCRUBBER, JsonPath, Scrubber
+from logfire.testing import TestExporter
 
 
 def chunks(format_string: str, kwargs: Mapping[str, Any]):
@@ -178,3 +181,50 @@ def test_internal_exception_formatting(caplog: pytest.LogCaptureFixture):
     assert len(caplog.records) == 1
     assert caplog.records[0].message.startswith('Caught an internal error in Logfire.')
     assert str(caplog.records[0].exc_info[1]) == 'bad scrubber'  # type: ignore
+
+
+@pytest.mark.skipif(sys.version_info[:2] == (3, 8), reason='fstring magic is only for 3.9+')
+@pytest.mark.anyio
+async def test_await_in_fstring(exporter: TestExporter):
+    """Test that logfire.info(f'{foo(await bar())}') evaluates the await expression and logs a warning."""
+
+    async def bar() -> str:
+        return 'content data'
+
+    def foo(x: str) -> str:
+        return x
+
+    with pytest.warns(FormattingFailedWarning) as warnings:
+        logfire.info(f'{foo(await bar())}')
+        [warning] = warnings
+        assert str(warning.message) == snapshot(
+            '\n'
+            '    Cannot evaluate await expression in f-string. Pre-evaluate the expression before logging.\n'
+            '    For example, change:\n'
+            '      logfire.info(f"{await get_value()}")\n'
+            '    To:\n'
+            '      value = await get_value()\n'
+            '      logfire.info(f"{value}")\n'
+            '    The problematic f-string value was: foo(await bar())'
+        )
+
+    assert exporter.exported_spans_as_dict() == snapshot(
+        [
+            {
+                'name': f'{foo(await bar())}',
+                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'parent': None,
+                'start_time': 1000000000,
+                'end_time': 1000000000,
+                'attributes': {
+                    'logfire.span_type': 'log',
+                    'logfire.level_num': 9,
+                    'logfire.msg_template': f'{foo(await bar())}',
+                    'logfire.msg': 'content data',
+                    'code.filepath': 'test_formatter.py',
+                    'code.lineno': 123,
+                    'code.function': 'test_await_in_fstring',
+                },
+            }
+        ]
+    )