Turn captions into list field

Signed-off-by: Christoph Auer <[email protected]>
DS4SD · Sep 23, 2024 · 425858a · 425858a
1 parent 0685709
commit 425858a
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/docling_core/types/__init__.py b/docling_core/types/__init__.py
@@ -23,3 +23,5 @@
 from docling_core.types.doc.document import ExportedCCSDocument as Document  # noqa
 from docling_core.types.gen.generic import Generic  # noqa
 from docling_core.types.rec.record import Record  # noqa
+
+from docling_core.types.experimental.document import DoclingDocument
diff --git a/docling_core/types/experimental/document.py b/docling_core/types/experimental/document.py
@@ -54,7 +54,7 @@ def from_dict_format(cls, data: Any) -> Any:
 
 
 class BaseTableData(BaseModel):  # TBD
-    table_cells: List[TableCell]
+    table_cells: List[TableCell] = []
     num_rows: int = 0
     num_cols: int = 0
 
@@ -184,7 +184,7 @@ class Section(TextItem):
 
 
 class FloatingItem(DocItem):
-    caption: Optional[RefItem] = None
+    captions: List[RefItem] = []
     references: List[RefItem] = []
     footnotes: List[RefItem] = []
     image: Optional[ImageRef] = None
@@ -333,7 +333,7 @@ def add_figure(
         if prov:
             fig_item.prov.append(prov)
         if caption:
-            fig_item.caption = caption.get_ref()
+            fig_item.captions.append(caption.get_ref())
 
         self.figures.append(fig_item)
         parent.children.append(RefItem(cref=cref))

diff --git a/test/data/experimental/dummy_doc.yaml b/test/data/experimental/dummy_doc.yaml
@@ -132,8 +132,8 @@ figures: # All figures...
     label: "figure"
     parent:
       $ref: "#/body"
-    caption:
-      $ref: "/texts/3"
+    captions:
+      - $ref: "/texts/3"
     data: # BaseFigureData Type
       classification: "illustration"
       confidence: 0.78