Support enums in format_as_xml (#3064)

Author: DouweM

pydantic_ai_slim/pydantic_ai/format_prompt.py

@@ -3,6 +3,7 @@
 from collections.abc import Iterable, Iterator, Mapping
 from dataclasses import asdict, dataclass, field, fields, is_dataclass
 from datetime import date
+from enum import Enum
 from typing import Any, Literal
 from xml.etree import ElementTree
 
@@ -26,8 +27,8 @@ def format_as_xml(
     This is useful since LLMs often find it easier to read semi-structured data (e.g. examples) as XML,
     rather than JSON etc.
 
-    Supports: `str`, `bytes`, `bytearray`, `bool`, `int`, `float`, `date`, `datetime`, `Mapping`,
-    `Iterable`, `dataclass`, and `BaseModel`.
+    Supports: `str`, `bytes`, `bytearray`, `bool`, `int`, `float`, `date`, `datetime`, `Enum`,
+    `Mapping`, `Iterable`, `dataclass`, and `BaseModel`.
 
     Args:
         obj: Python Object to serialize to XML.
@@ -101,7 +102,7 @@ def _to_xml(self, value: Any, path: str, tag: str | None = None) -> ElementTree.
             element.text = value
         elif isinstance(value, bytes | bytearray):
             element.text = value.decode(errors='ignore')
-        elif isinstance(value, bool | int | float):
+        elif isinstance(value, bool | int | float | Enum):
             element.text = str(value)
         elif isinstance(value, date):
             element.text = value.isoformat()

tests/test_format_as_xml.py

@@ -2,6 +2,7 @@
 
 from dataclasses import dataclass, field
 from datetime import date, datetime
+from enum import Enum
 from typing import Any
 
 import pytest
@@ -24,6 +25,16 @@ class ExamplePydanticModel(BaseModel):
     age: int
 
 
+class ExampleEnum(Enum):
+    FOO = 1
+    BAR = 2
+
+
+class ExampleStrEnum(str, Enum):
+    FOO = 'foo'
+    BAR = 'bar'
+
+
 class ExamplePydanticFields(BaseModel):
     name: str = Field(description="The person's name")
     age: int = Field(description='Years', title='Age', default=18)
@@ -43,6 +54,8 @@ def location(self) -> str | None:
         pytest.param('a string', snapshot('<examples>a string</examples>'), id='string'),
         pytest.param(42, snapshot('<examples>42</examples>'), id='int'),
         pytest.param(None, snapshot('<examples>null</examples>'), id='null'),
+        pytest.param(ExampleEnum.FOO, snapshot('<examples>ExampleEnum.FOO</examples>'), id='enum'),
+        pytest.param(ExampleStrEnum.FOO, snapshot('<examples>foo</examples>'), id='str enum'),
         pytest.param(
             ExampleDataclass(name='John', age=42),
             snapshot("""\