Fix Root Model Not Rendering As Expected (#1396)

Author: dhruv-ahuja

docs/reference/self-hosted/local-quickstart.md

@@ -11,7 +11,7 @@ For a production setup, including detailed configuration and prerequisites, plea
 
 Before deploying, you will need the following:
 
-- A Logfire Access Key, you'll need to get in contact with [[email protected]](mailto:[email protected]) to get one. **Remember you need to be on a trail for self-hosted enterprise Logfire to run logfire locally.**
+- A Logfire Access Key, you'll need to get in contact with [[email protected]](mailto:[email protected]) to get one. **Remember you need to be on a trial for self-hosted enterprise Logfire to run logfire locally.**
 - A local Kubernetes cluster, we will be using [Kind](https://kind.sigs.k8s.io/) in this example.
 - [Helm](https://helm.sh) CLI installed.
 - (Optional) [Tilt](https://tilt.dev/), as we will provide an optional convenience `Tiltfile` to automate the setup.

logfire/_internal/json_encoder.py

@@ -155,6 +155,10 @@ def _pydantic_model_encoder(o: Any, seen: set[int]) -> JsonValue:
     import pydantic
 
     assert isinstance(o, pydantic.BaseModel)
+
+    if isinstance(o, pydantic.RootModel):
+        return to_json_value(o.root, seen)  # type: ignore
+
     try:
         dump = o.model_dump()
     except AttributeError:  # pragma: no cover

logfire/_internal/json_formatter.py

@@ -19,6 +19,9 @@ def __init__(self, *, indent: int):
         self._indent_step = indent
         self._newlines = indent != 0
         self._data_type_map: dict[DataType, Callable[[int, Any, JSONSchema | None], None]] = {
+            'str': self._format_string,
+            'int': self._format_number,
+            'float': self._format_number,
             'PydanticModel': partial(self._format_items, '(', '=', ')', False),
             'dataclass': partial(self._format_items, '(', '=', ')', False),
             'Mapping': partial(self._format_items, '({', ': ', '})', True),
@@ -204,6 +207,38 @@ def _format_bytes(self, _indent_current: int, value: Any, schema: JSONSchema | N
         output = f'{cls}({value})' if cls else value
         self._stream.write(output)
 
+    def _format_string(self, _indent_current: int, value: Any, schema: JSONSchema | None) -> None:
+        """Format string value.
+
+        Examples:
+            >>> value = 'hello'
+            >>> schema = {'type': 'string', 'x-python-datatype': 'str'}
+            >>> _format_string(0, value, schema)
+            "hello"
+            >>> schema = {'type': 'string', 'x-python-datatype': 'str', 'title': 'MyString'}
+            >>> _format_string(0, value, schema)
+            MyString("hello")
+        """
+        cls = schema and schema.get('title')
+        output = f'{cls}({repr(value)})' if cls else repr(value)
+        self._stream.write(output)
+
+    def _format_number(self, _indent_current: int, value: Any, schema: JSONSchema | None) -> None:
+        """Format number value. Supports both integer and float types.
+
+        Examples:
+            >>> value = 42
+            >>> schema = {'type': 'integer', 'x-python-datatype': 'int'}
+            >>> _format_number(0, value, schema)
+            42
+            >>> schema = {'type': 'number', 'x-python-datatype': 'float', 'title': 'MyNumber'}
+            >>> _format_number(0, value, schema)
+            MyNumber(42)
+        """
+        cls = schema and schema.get('title')
+        output = f'{cls}({value})' if cls else str(value)
+        self._stream.write(output)
+
     def _format_table(
         self, columns: list[Any], indices: list[Any], rows: list[Any], real_column_count: int, real_row_count: int
     ) -> None:

logfire/_internal/json_schema.py

@@ -72,6 +72,7 @@ def type_to_schema() -> dict[type[Any], JsonDict | Callable[[Any, set[int]], Jso
                 pydantic.SecretBytes: {'type': 'string', 'x-python-datatype': 'SecretBytes'},
                 pydantic.AnyUrl: {'type': 'string', 'x-python-datatype': 'AnyUrl'},
                 pydantic.BaseModel: _pydantic_model_schema,
+                pydantic.RootModel: _pydantic_root_model_schema,
             }
         )
 
@@ -111,6 +112,7 @@ def create_json_schema(obj: Any, seen: set[int]) -> JsonDict:
     try:
         # cover common types first before calling `type_to_schema` to avoid the overhead of imports if not necessary
         obj_type = obj.__class__
+
         if obj_type in {str, int, bool, float}:
             return {}
 
@@ -282,10 +284,49 @@ def _exception_schema(obj: Exception, _seen: set[int]) -> JsonDict:
     return {'type': 'object', 'title': obj.__class__.__name__, 'x-python-datatype': 'Exception'}
 
 
+def _pydantic_root_model_schema(obj: Any, seen: set[int]) -> JsonDict:
+    import pydantic
+
+    assert isinstance(obj, pydantic.RootModel)
+
+    root = obj.root  # type: ignore
+
+    if isinstance(root, type(None)):
+        return {'type': 'null'}
+
+    schema: JsonDict = {}
+
+    # return a complex schema to ensure JSON parsing for simple objects inside RootModel since they get an
+    # extra layer of JSON encoding and to handle subclass representations correctly
+    primitive_types = (str, bool, int, float)
+    if isinstance(root, primitive_types):
+        datatype = None
+        if isinstance(root, str):
+            type_ = 'string'
+            datatype = 'str'
+        elif isinstance(root, bool):
+            type_ = 'boolean'
+        elif isinstance(root, int):
+            type_ = 'integer'
+            datatype = 'int'
+        else:
+            type_ = 'number'
+            datatype = 'float'
+
+        schema = {'type': type_, 'x-python-datatype': datatype}
+        if root.__class__ not in primitive_types:
+            schema['title'] = root.__class__.__name__
+
+        return schema
+
+    return create_json_schema(obj.root, seen)  # type: ignore
+
+
 def _pydantic_model_schema(obj: Any, seen: set[int]) -> JsonDict:
     import pydantic
 
     assert isinstance(obj, pydantic.BaseModel)
+
     try:
         fields = type(obj).model_fields
         extra = obj.model_extra or {}

logfire/_internal/json_types.py

@@ -11,6 +11,11 @@
     'Decimal',
     'UUID',
     'Enum',
+    # string
+    'str',
+    # number
+    'int',
+    'float',
     # bytes
     'bytes',
     # temporal types

tests/test_console_exporter.py

@@ -8,6 +8,7 @@
 import re
 import sys
 from datetime import datetime
+from enum import Enum
 from typing import Any
 from unittest import mock
 
@@ -31,6 +32,15 @@
 from logfire.testing import TestExporter
 from tests.utils import ReadableSpanModel, SpanContextModel, exported_spans_as_models
 
+if sys.version_info >= (3, 11):  # pragma: no branch
+    from enum import IntEnum, StrEnum
+else:  # pragma: no cover
+
+    class StrEnum(str, Enum): ...
+
+    class IntEnum(int, Enum): ...
+
+
 tracer = trace.get_tracer('test')
 
 NANOSECONDS_PER_SECOND = int(1e9)
@@ -1019,3 +1029,84 @@ def test_console_exporter_list_data_with_object_schema_mismatch(capsys: pytest.C
             "│ bar={'name': 'Alice', 'age': 30}",
         ]
     )
+
+
+def test_console_exporter_log_pydantic_root_model(capsys: pytest.CaptureFixture[str]) -> None:
+    from pydantic import BaseModel, RootModel
+
+    logfire.configure(
+        send_to_logfire=False,
+        console=ConsoleOptions(verbose=True, colors='never', include_timestamps=False),
+    )
+
+    class Model(BaseModel):
+        name: str
+
+    class Color(StrEnum):
+        red = 'RED'
+
+    class Order(IntEnum):
+        one = 1
+
+    RootWithModel = RootModel[Model]
+    RootWithStr = RootModel[str]
+    RootWithInt = RootModel[int]
+    RootWithFloat = RootModel[float]
+    RootWithBool = RootModel[bool]
+    RootWithNone = RootModel[None]
+    # enums (which are subclasses of their base types)
+    RootWithColor = RootModel[Color]
+    RootWithOrder = RootModel[Order]
+
+    model = Model(name='with_model')
+    root_with_model = RootWithModel(root=model)
+    root_with_str = RootWithStr('with_str')
+    root_with_int = RootWithInt(-150)
+    root_with_float = RootWithFloat(2.0)
+    root_with_bool = RootWithBool(False)
+    root_with_none = RootWithNone(None)
+    root_with_color = RootWithColor(Color.red)
+    root_with_order = RootWithOrder(Order.one)
+
+    logfire.info(
+        'hi',
+        with_model=root_with_model,
+        with_str=root_with_str,
+        with_str_inner=root_with_str.root,
+        with_int=root_with_int,
+        with_int_inner=root_with_int.root,
+        with_float=root_with_float,
+        with_float_inner=root_with_float.root,
+        with_bool=root_with_bool,
+        with_bool_inner=root_with_bool.root,
+        with_none=root_with_none,
+        with_none_inner=root_with_none.root,
+        with_color=root_with_color,
+        with_color_inner=root_with_color.root,
+        with_order=root_with_order,
+        with_order_inner=root_with_order.root,
+    )
+
+    assert capsys.readouterr().out.splitlines() == snapshot(
+        [
+            'hi',
+            IsStr(),
+            '│ with_model=Model(',
+            "│                name='with_model',",
+            '│            )',
+            "│ with_str='with_str'",
+            "│ with_str_inner='with_str'",
+            '│ with_int=-150',
+            '│ with_int_inner=-150',
+            '│ with_float=2.0',
+            '│ with_float_inner=2.0',
+            '│ with_bool=False',
+            '│ with_bool_inner=False',
+            '│ with_none=None',
+            '│ with_none_inner=None',
+            "│ with_color=Color('RED')",
+            "│ with_color_inner=Color('RED')",
+            '│ with_order=Order(1)',
+            '│ with_order_inner=Order(1)',
+        ]
+    )