feat(formatter): check if expression contains await

Author: Sbargaoui

logfire/_internal/ast_utils.py

@@ -106,7 +106,7 @@ def rewrite_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef, qualnam
 
         return ast.fix_missing_locations(
             ast.copy_location(
-                type(node)(  # type: ignore
+                type(node)(
                     name=node.name,
                     args=node.args,
                     body=new_body,

logfire/_internal/formatter.py

@@ -354,6 +354,21 @@ 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
+    try:
+        parsed = ast.parse(source, mode='eval')
+        for sub_node in ast.walk(parsed):
+            if isinstance(sub_node, ast.Await):
+                # Signal that we can't compile this
+                raise KnownFormattingError(
+                    f'Cannot evaluate await expression in f-string: {source}. Pre-evaluate the expression before logging.'
+                )
+    except SyntaxError:
+        # For other syntax errors during parsing, let the compile step handle it
+        # as it might provide better error messages
+        pass
+
     value_code = compile(source, '<fvalue1>', 'eval')
     expr = ast.Expression(
         ast.JoinedStr(

tests/test_formatter.py

@@ -178,3 +178,32 @@ 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
+
+
+def test_await_in_fstring():
+    """Test that await expressions in f-strings are rejected with a clear error message."""
+    # Create a simple AST node with an Await expression
+    import ast
+
+    # Create a FormattedValue node containing an Await expression
+    await_expr = ast.Await(value=ast.Name(id='coro', ctx=ast.Load()))
+    node = ast.FormattedValue(value=await_expr, conversion=-1, format_spec=None)
+
+    # Create a minimal mock of executing.Source
+    class MockSource:
+        text = "f'{await coro}'"
+
+        @staticmethod
+        def get_text(node: ast.AST) -> str:
+            return 'await coro'
+
+    mock_source = MockSource()
+
+    # Test that using compile_formatted_value raises the right error
+    from logfire._internal.formatter import KnownFormattingError, compile_formatted_value
+
+    with pytest.raises(KnownFormattingError) as exc_info:
+        compile_formatted_value(node, mock_source)
+
+    assert 'Cannot evaluate await expression in f-string' in str(exc_info.value)
+    assert 'Pre-evaluate the expression before logging' in str(exc_info.value)