🐛 Fix verbose console formatting for enum types (#1073) (#1096)

Co-authored-by: rahul-netomi <[email protected]>

Author: sbhrwlr

logfire/_internal/json_formatter.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import io
-from datetime import timedelta
+from datetime import date, datetime, time, timedelta
 from functools import partial
 from typing import Any, Callable, cast
 
@@ -30,9 +30,9 @@ def __init__(self, *, indent: int):
             'generator': partial(self._format_list_like, 'generator((', '))'),
             'bytes': self._format_bytes,
             'Decimal': partial(self._write, 'Decimal(', ')', True),
-            'date': partial(self._write, 'date(', ')', True),
-            'datetime': partial(self._write, 'datetime(', ')', True),
-            'time': partial(self._write, 'time(', ')', True),
+            'date': self._format_date,
+            'datetime': self._format_datetime,
+            'time': self._format_time,
             'timedelta': self._format_timedelta,
             'Enum': partial(self._write, '(', ')', True),
             'IPv4Address': partial(self._write, 'IPv4Address(', ')', True),
@@ -273,6 +273,15 @@ def _format_data_frame(self, _indent_current: int, value: list[Any], schema: JSO
             real_row_count=schema.get('x-row-count', 0),
         )
 
+    def _format_date(self, _indent_current: int, value: Any, schema: JSONSchema | None) -> None:
+        self._write('', '', True, 0, date.fromisoformat(value) if isinstance(value, str) else value, schema)
+
+    def _format_datetime(self, _indent_current: int, value: Any, schema: JSONSchema | None) -> None:
+        self._write('', '', True, 0, datetime.fromisoformat(value) if isinstance(value, str) else value, schema)
+
+    def _format_time(self, _indent_current: int, value: Any, schema: JSONSchema | None) -> None:
+        self._write('', '', True, 0, time.fromisoformat(value) if isinstance(value, str) else value, schema)
+
 
 json_args_value_formatter = JsonArgsValueFormatter(indent=4)
 json_args_value_formatter_compact = JsonArgsValueFormatter(indent=0)

logfire/_internal/json_schema.py

@@ -43,10 +43,10 @@ def type_to_schema() -> dict[type[Any], JsonDict | Callable[[Any, set[int]], Jso
         bytes: _bytes_schema,
         bytearray: _bytearray_schema,
         Enum: _enum_schema,
-        Decimal: {'type': 'string', 'format': 'decimal'},
-        datetime.datetime: {'type': 'string', 'format': 'date-time'},
-        datetime.date: {'type': 'string', 'format': 'date'},
-        datetime.time: {'type': 'string', 'format': 'time'},
+        Decimal: {'type': 'string', 'format': 'decimal', 'x-python-datatype': 'Decimal'},
+        datetime.datetime: {'type': 'string', 'format': 'date-time', 'x-python-datatype': 'datetime'},
+        datetime.date: {'type': 'string', 'format': 'date', 'x-python-datatype': 'date'},
+        datetime.time: {'type': 'string', 'format': 'time', 'x-python-datatype': 'time'},
         datetime.timedelta: {'type': 'string', 'x-python-datatype': 'timedelta'},
         GeneratorType: _generator_schema,
         IPv4Address: {'type': 'string', 'format': 'ipv4'},

logfire/_internal/main.py

@@ -6,6 +6,7 @@
 import sys
 import warnings
 from contextvars import Token
+from enum import Enum
 from functools import cached_property
 from time import time
 from typing import (
@@ -2450,7 +2451,9 @@ def prepare_otlp_attributes(attributes: dict[str, Any]) -> dict[str, otel_types.
 
 def prepare_otlp_attribute(value: Any) -> otel_types.AttributeValue:
     """Convert a user attribute to an OpenTelemetry compatible type."""
-    if isinstance(value, int):
+    if isinstance(value, Enum):
+        return logfire_json_dumps(value)
+    elif isinstance(value, int):
         if value > OTLP_MAX_INT_SIZE:
             warnings.warn(
                 f'Integer value {value} is larger than the maximum OTLP integer size of {OTLP_MAX_INT_SIZE} (64-bits), '

tests/test_console_exporter.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import decimal
+import enum
 import io
 import sys
 from datetime import datetime
@@ -943,19 +944,40 @@ def test_other_json_schema_types(capsys: pytest.CaptureFixture[str]) -> None:
         console=ConsoleOptions(verbose=True, colors='never', include_timestamps=False),
     )
 
+    class MyEnum(enum.Enum):
+        """Enum with string values."""
+
+        ABC = 'abc'
+
+    class MyStrEnum(str, enum.Enum):
+        """String-based Enum."""
+
+        STR = 'str_val'
+
+    class MyIntEnum(int, enum.Enum):
+        """Integer-based Enum."""
+
+        INT = 1
+
     logfire.info(
         'hi',
         d=datetime(2020, 12, 31, 12, 34, 56),
         x=None,
         v=decimal.Decimal('1.0'),
+        e=MyEnum.ABC,
+        se=MyStrEnum.STR,
+        ie=MyIntEnum.INT,
     )
 
     assert capsys.readouterr().out.splitlines() == snapshot(
         [
             'hi',
             IsStr(),
-            '│ d=2020-12-31T12:34:56',
+            '│ d=datetime.datetime(2020, 12, 31, 12, 34, 56)',
             '│ x=None',
-            '│ v=1.0',
+            "│ v=Decimal('1.0')",
+            "│ e=MyEnum('abc')",
+            "│ se=MyStrEnum('str_val')",
+            '│ ie=MyIntEnum(1)',
         ]
     )

tests/test_json_args.py

@@ -317,28 +317,28 @@ class StrSubclass(str):
             Decimal('1.7'),
             '1.7',
             '"1.7"',
-            {'type': 'string', 'format': 'decimal'},
+            {'type': 'string', 'format': 'decimal', 'x-python-datatype': 'Decimal'},
             id='decimal',
         ),
         pytest.param(
             date(2023, 1, 1),
             '2023-01-01',
             '"2023-01-01"',
-            {'type': 'string', 'format': 'date'},
+            {'type': 'string', 'format': 'date', 'x-python-datatype': 'date'},
             id='date',
         ),
         pytest.param(
             datetime(2023, 1, 1, 10, 10),
             '2023-01-01 10:10:00',
             '"2023-01-01T10:10:00"',
-            {'type': 'string', 'format': 'date-time'},
+            {'type': 'string', 'format': 'date-time', 'x-python-datatype': 'datetime'},
             id='datetime',
         ),
         pytest.param(
             time(12, 10),
             '12:10:00',
             '"12:10:00"',
-            {'type': 'string', 'format': 'time'},
+            {'type': 'string', 'format': 'time', 'x-python-datatype': 'time'},
             id='time',
         ),
         pytest.param(
@@ -973,7 +973,7 @@ class MyPydanticDataclass:
                 '...'
                 'og_non_scalar_complex_args.<locals>.MyPydanticDataclass(p=20)}'
             ),
-            'logfire.json_schema': '{"type":"object","properties":{"a":{},"complex_list":{"type":"array","prefixItems":[{},{},{"type":"object","title":"MyModel","x-python-datatype":"PydanticModel","properties":{"y":{"type":"string","format":"date-time"}}},{"type":"object","title":"MyDataclass","x-python-datatype":"dataclass"},{"type":"object","title":"MyPydanticDataclass","x-python-datatype":"dataclass"}]},"complex_dict":{"type":"object","properties":{"model":{"type":"object","title":"MyModel","x-python-datatype":"PydanticModel","properties":{"y":{"type":"string","format":"date-time"}}},"dataclass":{"type":"object","title":"MyDataclass","x-python-datatype":"dataclass"},"pydantic_dataclass":{"type":"object","title":"MyPydanticDataclass","x-python-datatype":"dataclass"}}}}}',
+            'logfire.json_schema': '{"type":"object","properties":{"a":{},"complex_list":{"type":"array","prefixItems":[{},{},{"type":"object","title":"MyModel","x-python-datatype":"PydanticModel","properties":{"y":{"type":"string","format":"date-time","x-python-datatype":"datetime"}}},{"type":"object","title":"MyDataclass","x-python-datatype":"dataclass"},{"type":"object","title":"MyPydanticDataclass","x-python-datatype":"dataclass"}]},"complex_dict":{"type":"object","properties":{"model":{"type":"object","title":"MyModel","x-python-datatype":"PydanticModel","properties":{"y":{"type":"string","format":"date-time","x-python-datatype":"datetime"}}},"dataclass":{"type":"object","title":"MyDataclass","x-python-datatype":"dataclass"},"pydantic_dataclass":{"type":"object","title":"MyPydanticDataclass","x-python-datatype":"dataclass"}}}}}',
             'code.filepath': 'test_json_args.py',
             'code.lineno': 123,
             'code.function': 'test_log_non_scalar_complex_args',

tests/test_json_args_formatting.py

@@ -148,19 +148,19 @@
         pytest.param(
             '2023-01-01',
             {'type': 'string', 'x-python-datatype': 'date'},
-            "date('2023-01-01')",
+            'datetime.date(2023, 1, 1)',
             id='date',
         ),
         pytest.param(
             '2023-01-01T10:10:00',
             {'type': 'string', 'x-python-datatype': 'datetime'},
-            "datetime('2023-01-01T10:10:00')",
+            'datetime.datetime(2023, 1, 1, 10, 10)',
             id='datetime',
         ),
         pytest.param(
             '12:10:00',
             {'type': 'string', 'x-python-datatype': 'time'},
-            "time('12:10:00')",
+            'datetime.time(12, 10)',
             id='time',
         ),
         pytest.param(
@@ -494,7 +494,7 @@ def test_nested_json_args_value_formatting():
     1,
     MyModel(
         x='x',
-        y=datetime('2023-01-01T00:00:00'),
+        y=datetime.datetime(2023, 1, 1, 0, 0),
     ),
     MyDataclass(
         t=10,
@@ -554,19 +554,19 @@ def test_nested_json_args_value_formatting():
         pytest.param(
             '2023-01-01',
             {'type': 'string', 'x-python-datatype': 'date'},
-            "date('2023-01-01')",
+            'datetime.date(2023, 1, 1)',
             id='date',
         ),
         pytest.param(
             '2023-01-01T10:10:00',
             {'type': 'string', 'x-python-datatype': 'datetime'},
-            "datetime('2023-01-01T10:10:00')",
+            'datetime.datetime(2023, 1, 1, 10, 10)',
             id='datetime',
         ),
         pytest.param(
             '12:10:00',
             {'type': 'string', 'x-python-datatype': 'time'},
-            "time('12:10:00')",
+            'datetime.time(12, 10)',
             id='time',
         ),
         pytest.param(