Use model class names as tags in format_as_xml and add option to include field titles and descriptions as attributes (#2313)

Co-authored-by: giacbrd <[email protected]>
Co-authored-by: Douwe Maan <[email protected]>

Author: 3 people

pydantic_ai_slim/pydantic_ai/format_prompt.py

@@ -1,22 +1,25 @@
 from __future__ import annotations as _annotations
 
 from collections.abc import Iterable, Iterator, Mapping
-from dataclasses import asdict, dataclass, is_dataclass
+from dataclasses import asdict, dataclass, field, fields, is_dataclass
 from datetime import date
-from typing import Any
+from typing import Any, Literal
 from xml.etree import ElementTree
 
 from pydantic import BaseModel
 
 __all__ = ('format_as_xml',)
 
+from pydantic.fields import ComputedFieldInfo, FieldInfo
+
 
 def format_as_xml(
     obj: Any,
     root_tag: str | None = None,
     item_tag: str = 'item',
     none_str: str = 'null',
     indent: str | None = '  ',
+    include_field_info: Literal['once'] | bool = False,
 ) -> str:
     """Format a Python object as XML.
 
@@ -33,6 +36,10 @@ def format_as_xml(
             for dataclasses and Pydantic models.
         none_str: String to use for `None` values.
         indent: Indentation string to use for pretty printing.
+        include_field_info: Whether to include attributes like Pydantic `Field` attributes and dataclasses `field()`
+            `metadata` as XML attributes. In both cases the allowed `Field` attributes and `field()` metadata keys are
+            `title` and `description`. If a field is repeated in the data (e.g. in a list) by setting `once`
+            the attributes are included only in the first occurrence of an XML element relative to the same field.
 
     Returns:
         XML representation of the object.
@@ -51,7 +58,12 @@ def format_as_xml(
     '''
     ```
     """
-    el = _ToXml(item_tag=item_tag, none_str=none_str).to_xml(obj, root_tag)
+    el = _ToXml(
+        data=obj,
+        item_tag=item_tag,
+        none_str=none_str,
+        include_field_info=include_field_info,
+    ).to_xml(root_tag)
     if root_tag is None and el.text is None:
         join = '' if indent is None else '\n'
         return join.join(_rootless_xml_elements(el, indent))
@@ -63,11 +75,26 @@ def format_as_xml(
 
 @dataclass
 class _ToXml:
+    data: Any
     item_tag: str
     none_str: str
-
-    def to_xml(self, value: Any, tag: str | None) -> ElementTree.Element:
-        element = ElementTree.Element(self.item_tag if tag is None else tag)
+    include_field_info: Literal['once'] | bool
+    # a map of Pydantic and dataclasses Field paths to their metadata:
+    # a field unique string representation and its class
+    _fields_info: dict[str, tuple[str, FieldInfo | ComputedFieldInfo]] = field(default_factory=dict)
+    # keep track of fields we have extracted attributes from
+    _included_fields: set[str] = field(default_factory=set)
+    # keep track of class names for dataclasses and Pydantic models, that occur in lists
+    _element_names: dict[str, str] = field(default_factory=dict)
+    # flag for parsing dataclasses and Pydantic models once
+    _is_info_extracted: bool = False
+    _FIELD_ATTRIBUTES = ('title', 'description')
+
+    def to_xml(self, tag: str | None = None) -> ElementTree.Element:
+        return self._to_xml(value=self.data, path='', tag=tag)
+
+    def _to_xml(self, value: Any, path: str, tag: str | None = None) -> ElementTree.Element:
+        element = self._create_element(self.item_tag if tag is None else tag, path)
         if value is None:
             element.text = self.none_str
         elif isinstance(value, str):
@@ -79,31 +106,96 @@ def to_xml(self, value: Any, tag: str | None) -> ElementTree.Element:
         elif isinstance(value, date):
             element.text = value.isoformat()
         elif isinstance(value, Mapping):
-            self._mapping_to_xml(element, value)  # pyright: ignore[reportUnknownArgumentType]
+            if tag is None and path in self._element_names:
+                element.tag = self._element_names[path]
+            self._mapping_to_xml(element, value, path)  # pyright: ignore[reportUnknownArgumentType]
         elif is_dataclass(value) and not isinstance(value, type):
+            self._init_structure_info()
             if tag is None:
-                element = ElementTree.Element(value.__class__.__name__)
-            dc_dict = asdict(value)
-            self._mapping_to_xml(element, dc_dict)
+                element.tag = value.__class__.__name__
+            self._mapping_to_xml(element, asdict(value), path)
         elif isinstance(value, BaseModel):
+            self._init_structure_info()
             if tag is None:
-                element = ElementTree.Element(value.__class__.__name__)
-            self._mapping_to_xml(element, value.model_dump(mode='python'))
+                element.tag = value.__class__.__name__
+            # by dumping the model we loose all metadata in nested data structures,
+            # but we have collected it when called _init_structure_info
+            self._mapping_to_xml(element, value.model_dump(), path)
         elif isinstance(value, Iterable):
-            for item in value:  # pyright: ignore[reportUnknownVariableType]
-                item_el = self.to_xml(item, None)
-                element.append(item_el)
+            for n, item in enumerate(value):  # pyright: ignore[reportUnknownVariableType,reportUnknownArgumentType]
+                element.append(self._to_xml(value=item, path=f'{path}.[{n}]' if path else f'[{n}]'))
         else:
             raise TypeError(f'Unsupported type for XML formatting: {type(value)}')
         return element
 
-    def _mapping_to_xml(self, element: ElementTree.Element, mapping: Mapping[Any, Any]) -> None:
+    def _create_element(self, tag: str, path: str) -> ElementTree.Element:
+        element = ElementTree.Element(tag)
+        if path in self._fields_info:
+            field_repr, field_info = self._fields_info[path]
+            if self.include_field_info and self.include_field_info != 'once' or field_repr not in self._included_fields:
+                field_attributes = self._extract_attributes(field_info)
+                for k, v in field_attributes.items():
+                    element.set(k, v)
+                self._included_fields.add(field_repr)
+        return element
+
+    def _init_structure_info(self):
+        """Create maps with all data information (fields info and class names), if not already created."""
+        if not self._is_info_extracted:
+            self._parse_data_structures(self.data)
+            self._is_info_extracted = True
+
+    def _mapping_to_xml(
+        self,
+        element: ElementTree.Element,
+        mapping: Mapping[Any, Any],
+        path: str = '',
+    ) -> None:
         for key, value in mapping.items():
             if isinstance(key, int):
                 key = str(key)
             elif not isinstance(key, str):
                 raise TypeError(f'Unsupported key type for XML formatting: {type(key)}, only str and int are allowed')
-            element.append(self.to_xml(value, key))
+            element.append(self._to_xml(value=value, path=f'{path}.{key}' if path else key, tag=key))
+
+    def _parse_data_structures(
+        self,
+        value: Any,
+        path: str = '',
+    ):
+        """Parse data structures as dataclasses or Pydantic models to extract element names and attributes."""
+        if value is None or isinstance(value, (str | int | float | date | bytearray | bytes | bool)):
+            return
+        elif isinstance(value, Mapping):
+            for k, v in value.items():  # pyright: ignore[reportUnknownVariableType]
+                self._parse_data_structures(v, f'{path}.{k}' if path else f'{k}')
+        elif is_dataclass(value) and not isinstance(value, type):
+            self._element_names[path] = value.__class__.__name__
+            for field in fields(value):
+                new_path = f'{path}.{field.name}' if path else field.name
+                if self.include_field_info and field.metadata:
+                    attributes = {k: v for k, v in field.metadata.items() if k in self._FIELD_ATTRIBUTES}
+                    if attributes:
+                        field_repr = f'{value.__class__.__name__}.{field.name}'
+                        self._fields_info[new_path] = (field_repr, FieldInfo(**attributes))
+                self._parse_data_structures(getattr(value, field.name), new_path)
+        elif isinstance(value, BaseModel):
+            self._element_names[path] = value.__class__.__name__
+            for model_fields in (value.__class__.model_fields, value.__class__.model_computed_fields):
+                for field, info in model_fields.items():
+                    new_path = f'{path}.{field}' if path else field
+                    if self.include_field_info and (isinstance(info, ComputedFieldInfo) or not info.exclude):
+                        field_repr = f'{value.__class__.__name__}.{field}'
+                        self._fields_info[new_path] = (field_repr, info)
+                    self._parse_data_structures(getattr(value, field), new_path)
+        elif isinstance(value, Iterable):
+            for n, item in enumerate(value):  # pyright: ignore[reportUnknownVariableType,reportUnknownArgumentType]
+                new_path = f'{path}.[{n}]' if path else f'[{n}]'
+                self._parse_data_structures(item, new_path)
+
+    @classmethod
+    def _extract_attributes(cls, info: FieldInfo | ComputedFieldInfo) -> dict[str, str]:
+        return {attr: str(value) for attr in cls._FIELD_ATTRIBUTES if (value := getattr(info, attr, None)) is not None}
 
 
 def _rootless_xml_elements(root: ElementTree.Element, indent: str | None) -> Iterator[str]: