chore: add str and int enum tests to cover type subclasses

Author: dhruv-ahuja

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)
@@ -1032,12 +1042,21 @@ def test_console_exporter_log_pydantic_root_model(capsys: pytest.CaptureFixture[
     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)
@@ -1046,6 +1065,8 @@ class Model(BaseModel):
     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',
@@ -1060,6 +1081,10 @@ class Model(BaseModel):
         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(
@@ -1079,5 +1104,9 @@ class Model(BaseModel):
             '│ 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)',
         ]
     )

tests/test_json_args.py

@@ -33,11 +33,13 @@
 from logfire.testing import TestExporter
 
 if sys.version_info >= (3, 11):  # pragma: no branch
-    from enum import StrEnum
+    from enum import IntEnum, StrEnum
 else:  # pragma: no cover
 
     class StrEnum(str, Enum): ...
 
+    class IntEnum(int, Enum): ...
+
 
 pandas.set_option('display.max_columns', 10)
 pandas.set_option('display.max_rows', 20)
@@ -1384,12 +1386,21 @@ def test_pydantic_root_model(exporter: TestExporter):
     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)
@@ -1398,6 +1409,8 @@ class Model(BaseModel):
     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',
@@ -1412,6 +1425,10 @@ class Model(BaseModel):
         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 exporter.exported_spans_as_dict() == [
@@ -1442,7 +1459,11 @@ class Model(BaseModel):
                 'with_bool_inner': False,
                 'with_none': 'null',
                 'with_none_inner': 'null',
-                'logfire.json_schema': '{"type":"object","properties":{"with_model":{"type":"object","title":"Model","x-python-datatype":"PydanticModel"},"with_str":{"type":"string","x-python-datatype":"string","title":"str"},"with_str_inner":{},"with_int":{"type":"integer","x-python-datatype":"number","title":"int"},"with_int_inner":{},"with_float":{"type":"number","x-python-datatype":"number","title":"float"},"with_float_inner":{},"with_bool":{"type":"boolean","x-python-datatype":null,"title":"bool"},"with_bool_inner":{},"with_none":{"type":"null"},"with_none_inner":{"type":"null"}}}',
+                'with_color': '"RED"',
+                'with_color_inner': '"RED"',
+                'with_order': '1',
+                'with_order_inner': '1',
+                'logfire.json_schema': '{"type":"object","properties":{"with_model":{"type":"object","title":"Model","x-python-datatype":"PydanticModel"},"with_str":{"type":"string","x-python-datatype":"string","title":"str"},"with_str_inner":{},"with_int":{"type":"integer","x-python-datatype":"number","title":"int"},"with_int_inner":{},"with_float":{"type":"number","x-python-datatype":"number","title":"float"},"with_float_inner":{},"with_bool":{"type":"boolean","x-python-datatype":null,"title":"bool"},"with_bool_inner":{},"with_none":{"type":"null"},"with_none_inner":{"type":"null"},"with_color":{"type":"string","x-python-datatype":"string","title":"Color"},"with_color_inner":{"type":"string","title":"Color","x-python-datatype":"Enum","enum":["RED"]},"with_order":{"type":"integer","x-python-datatype":"number","title":"Order"},"with_order_inner":{"type":"integer","title":"Order","x-python-datatype":"Enum","enum":[1]}}}',
             },
         }
     ]