feat: enable precision control in float serialization (#352)

Signed-off-by: Panos Vagenas <[email protected]>

Author: vagenas

docling_core/types/doc/base.py

@@ -3,7 +3,7 @@
 from enum import Enum
 from typing import List, Tuple
 
-from pydantic import BaseModel
+from pydantic import BaseModel, FieldSerializationInfo, field_serializer
 
 
 class ImageRefMode(str, Enum):
@@ -21,12 +21,28 @@ class CoordOrigin(str, Enum):
     BOTTOMLEFT = "BOTTOMLEFT"
 
 
+_CTX_COORD_PREC = "coord_prec"
+
+
+def _serialize_precision(
+    value: float, info: FieldSerializationInfo, ctx_key: str
+) -> float:
+    precision = info.context.get(ctx_key) if info.context else None
+    if isinstance(precision, int):
+        return round(value, precision)
+    return value
+
+
 class Size(BaseModel):
     """Size."""
 
     width: float = 0.0
     height: float = 0.0
 
+    @field_serializer("width", "height")
+    def _serialize(self, value: float, info: FieldSerializationInfo) -> float:
+        return _serialize_precision(value, info, _CTX_COORD_PREC)
+
     def as_tuple(self):
         """as_tuple."""
         return (self.width, self.height)
@@ -52,6 +68,10 @@ def height(self):
         """height."""
         return abs(self.t - self.b)
 
+    @field_serializer("l", "t", "r", "b")
+    def _serialize(self, value: float, info: FieldSerializationInfo) -> float:
+        return _serialize_precision(value, info, _CTX_COORD_PREC)
+
     def resize_by_scale(self, x_scale: float, y_scale: float):
         """resize_by_scale."""
         return BoundingBox(

docling_core/types/doc/document.py

@@ -26,8 +26,10 @@
     BaseModel,
     ConfigDict,
     Field,
+    FieldSerializationInfo,
     StringConstraints,
     computed_field,
+    field_serializer,
     field_validator,
     model_validator,
     validate_call,
@@ -38,7 +40,12 @@
 from docling_core.search.package import VERSION_PATTERN
 from docling_core.types.base import _JSON_POINTER_REGEX
 from docling_core.types.doc import BoundingBox, Size
-from docling_core.types.doc.base import CoordOrigin, ImageRefMode
+from docling_core.types.doc.base import (
+    _CTX_COORD_PREC,
+    CoordOrigin,
+    ImageRefMode,
+    _serialize_precision,
+)
 from docling_core.types.doc.labels import (
     CodeLanguageLabel,
     DocItemLabel,
@@ -85,6 +92,8 @@
     ]
 )
 
+_CTX_CONFID_PREC = "confid_prec"
+
 
 class BaseAnnotation(BaseModel):
     """Base class for all annotation types."""
@@ -98,6 +107,10 @@ class PictureClassificationClass(BaseModel):
     class_name: str
     confidence: float
 
+    @field_serializer("confidence")
+    def _serialize(self, value: float, info: FieldSerializationInfo) -> float:
+        return _serialize_precision(value, info, _CTX_CONFID_PREC)
+
 
 class PictureClassificationData(BaseAnnotation):
     """PictureClassificationData."""
@@ -125,6 +138,10 @@ class PictureMoleculeData(BaseAnnotation):
     segmentation: List[Tuple[float, float]]
     provenance: str
 
+    @field_serializer("confidence")
+    def _serialize(self, value: float, info: FieldSerializationInfo) -> float:
+        return _serialize_precision(value, info, _CTX_CONFID_PREC)
+
 
 class MiscAnnotation(BaseAnnotation):
     """MiscAnnotation."""
@@ -3048,6 +3065,8 @@ def save_as_json(
         artifacts_dir: Optional[Path] = None,
         image_mode: ImageRefMode = ImageRefMode.EMBEDDED,
         indent: int = 2,
+        coord_precision: Optional[int] = None,
+        confid_precision: Optional[int] = None,
     ):
         """Save as json."""
         if isinstance(filename, str):
@@ -3061,7 +3080,9 @@ def save_as_json(
             artifacts_dir, image_mode, reference_path=reference_path
         )
 
-        out = new_doc.export_to_dict()
+        out = new_doc.export_to_dict(
+            coord_precision=coord_precision, confid_precision=confid_precision
+        )
         with open(filename, "w", encoding="utf-8") as fw:
             json.dump(out, fw, indent=indent)
 
@@ -3087,6 +3108,8 @@ def save_as_yaml(
         artifacts_dir: Optional[Path] = None,
         image_mode: ImageRefMode = ImageRefMode.EMBEDDED,
         default_flow_style: bool = False,
+        coord_precision: Optional[int] = None,
+        confid_precision: Optional[int] = None,
     ):
         """Save as yaml."""
         if isinstance(filename, str):
@@ -3100,7 +3123,9 @@ def save_as_yaml(
             artifacts_dir, image_mode, reference_path=reference_path
         )
 
-        out = new_doc.export_to_dict()
+        out = new_doc.export_to_dict(
+            coord_precision=coord_precision, confid_precision=confid_precision
+        )
         with open(filename, "w", encoding="utf-8") as fw:
             yaml.dump(out, fw, default_flow_style=default_flow_style)
 
@@ -3125,9 +3150,18 @@ def export_to_dict(
         mode: str = "json",
         by_alias: bool = True,
         exclude_none: bool = True,
+        coord_precision: Optional[int] = None,
+        confid_precision: Optional[int] = None,
     ) -> Dict[str, Any]:
         """Export to dict."""
-        out = self.model_dump(mode=mode, by_alias=by_alias, exclude_none=exclude_none)
+        context = {}
+        if coord_precision is not None:
+            context[_CTX_COORD_PREC] = coord_precision
+        if confid_precision is not None:
+            context[_CTX_CONFID_PREC] = confid_precision
+        out = self.model_dump(
+            mode=mode, by_alias=by_alias, exclude_none=exclude_none, context=context
+        )
 
         return out
 

docling_core/types/doc/page.py

@@ -25,9 +25,21 @@
 from PIL import Image as PILImage
 from PIL import ImageColor, ImageDraw, ImageFont
 from PIL.ImageFont import FreeTypeFont
-from pydantic import AnyUrl, BaseModel, Field, model_validator
+from pydantic import (
+    AnyUrl,
+    BaseModel,
+    Field,
+    FieldSerializationInfo,
+    field_serializer,
+    model_validator,
+)
 
-from docling_core.types.doc.base import BoundingBox, CoordOrigin
+from docling_core.types.doc.base import (
+    _CTX_COORD_PREC,
+    BoundingBox,
+    CoordOrigin,
+    _serialize_precision,
+)
 from docling_core.types.doc.document import ImageRef
 
 _logger = logging.getLogger(__name__)
@@ -105,6 +117,10 @@ class BoundingRectangle(BaseModel):
 
     coord_origin: CoordOrigin = CoordOrigin.BOTTOMLEFT
 
+    @field_serializer("r_x0", "r_y0", "r_x1", "r_y1", "r_x2", "r_y2", "r_x3", "r_y3")
+    def _serialize(self, value: float, info: FieldSerializationInfo) -> float:
+        return _serialize_precision(value, info, _CTX_COORD_PREC)
+
     @property
     def width(self) -> float:
         """Calculate the width of the rectangle."""