fix: use same base type for all components (#10)

Signed-off-by: Michele Dolfi <[email protected]>

Author: dolfim-ibm

docling_core/types/doc/base.py

@@ -128,37 +128,28 @@ class GlmTableCell(TableCell):
     )
 
 
-class Table(AliasModel):
-    """Table."""
+class BaseCell(AliasModel):
+    """Base cell."""
 
-    num_cols: int = Field(alias="#-cols")
-    num_rows: int = Field(alias="#-rows")
     bounding_box: Optional[BoundingBoxContainer] = Field(
         default=None, alias="bounding-box", json_schema_extra=es_field(suppress=True)
     )
-    data: Optional[list[list[Union[GlmTableCell, TableCell]]]] = None
-    model: Optional[str] = None
     prov: Optional[list[Prov]] = None
     text: Optional[str] = Field(
         default=None, json_schema_extra=es_field(term_vector="with_positions_offsets")
     )
     obj_type: str = Field(
-        alias="type",
-        json_schema_extra=es_field(type="keyword", ignore_above=8191),
+        alias="type", json_schema_extra=es_field(type="keyword", ignore_above=8191)
     )
 
 
-class BaseCell(AliasModel):
-    """Base cell."""
+class Table(BaseCell):
+    """Table."""
 
-    bounding_box: Optional[BoundingBoxContainer] = Field(
-        default=None, alias="bounding-box", json_schema_extra=es_field(suppress=True)
-    )
-    prov: Optional[list[Prov]] = None
-    text: Optional[str] = None
-    obj_type: str = Field(
-        alias="type", json_schema_extra=es_field(type="keyword", ignore_above=8191)
-    )
+    num_cols: int = Field(alias="#-cols")
+    num_rows: int = Field(alias="#-rows")
+    data: Optional[list[list[Union[GlmTableCell, TableCell]]]] = None
+    model: Optional[str] = None
 
 
 class BaseText(AliasModel):

docling_core/types/doc/document.py

@@ -393,14 +393,28 @@ def from_dict(cls, data):
 
         return data
 
-    def _resolve_ref(self, item: Ref) -> Optional[Table]:
-        """Return the resolved reference in case of table reference, otherwise None."""
-        result: Optional[Table] = None
+    def _resolve_ref(self, item: Ref) -> Optional[Union[BaseCell, BaseText]]:
+        """Return the resolved reference.
 
-        # NOTE: currently only resolves table refs & makes assumptions on ref parts
+        Resolved the Ref object within the document.
+        If the object is not found, None is returned.
+        """
+        result: Optional[Union[BaseCell, BaseText]] = None
+
+        # NOTE: currently only resolves refs explicitely, such that we can make
+        # assumptions on ref parts
         if item.obj_type == "table" and self.tables:
             parts = item.ref.split("/")
             result = self.tables[int(parts[2])]
+        elif item.obj_type == "figure" and self.figures:
+            parts = item.ref.split("/")
+            result = self.figures[int(parts[2])]
+        elif item.obj_type == "equation" and self.equations:
+            parts = item.ref.split("/")
+            result = self.equations[int(parts[2])]
+        elif item.obj_type == "footnote" and self.footnotes:
+            parts = item.ref.split("/")
+            result = self.footnotes[int(parts[2])]
 
         return result
 

docs/Document.json

@@ -236,7 +236,8 @@
             }
           ],
           "default": null,
-          "title": "Text"
+          "title": "Text",
+          "x-es-term_vector": "with_positions_offsets"
         },
         "type": {
           "title": "Type",
@@ -1571,14 +1572,6 @@
     "Table": {
       "description": "Table.",
       "properties": {
-        "#-cols": {
-          "title": "#-Cols",
-          "type": "integer"
-        },
-        "#-rows": {
-          "title": "#-Rows",
-          "type": "integer"
-        },
         "bounding-box": {
           "anyOf": [
             {
@@ -1591,21 +1584,11 @@
           "default": null,
           "x-es-suppress": true
         },
-        "data": {
+        "prov": {
           "anyOf": [
             {
               "items": {
-                "items": {
-                  "anyOf": [
-                    {
-                      "$ref": "#/$defs/GlmTableCell"
-                    },
-                    {
-                      "$ref": "#/$defs/TableCell"
-                    }
-                  ]
-                },
-                "type": "array"
+                "$ref": "#/$defs/Prov"
               },
               "type": "array"
             },
@@ -1614,9 +1597,9 @@
             }
           ],
           "default": null,
-          "title": "Data"
+          "title": "Prov"
         },
-        "model": {
+        "text": {
           "anyOf": [
             {
               "type": "string"
@@ -1626,13 +1609,38 @@
             }
           ],
           "default": null,
-          "title": "Model"
+          "title": "Text",
+          "x-es-term_vector": "with_positions_offsets"
         },
-        "prov": {
+        "type": {
+          "title": "Type",
+          "type": "string",
+          "x-es-ignore_above": 8191,
+          "x-es-type": "keyword"
+        },
+        "#-cols": {
+          "title": "#-Cols",
+          "type": "integer"
+        },
+        "#-rows": {
+          "title": "#-Rows",
+          "type": "integer"
+        },
+        "data": {
           "anyOf": [
             {
               "items": {
-                "$ref": "#/$defs/Prov"
+                "items": {
+                  "anyOf": [
+                    {
+                      "$ref": "#/$defs/GlmTableCell"
+                    },
+                    {
+                      "$ref": "#/$defs/TableCell"
+                    }
+                  ]
+                },
+                "type": "array"
               },
               "type": "array"
             },
@@ -1641,9 +1649,9 @@
             }
           ],
           "default": null,
-          "title": "Prov"
+          "title": "Data"
         },
-        "text": {
+        "model": {
           "anyOf": [
             {
               "type": "string"
@@ -1653,20 +1661,13 @@
             }
           ],
           "default": null,
-          "title": "Text",
-          "x-es-term_vector": "with_positions_offsets"
-        },
-        "type": {
-          "title": "Type",
-          "type": "string",
-          "x-es-ignore_above": 8191,
-          "x-es-type": "keyword"
+          "title": "Model"
         }
       },
       "required": [
+        "type",
         "#-cols",
-        "#-rows",
-        "type"
+        "#-rows"
       ],
       "title": "Table",
       "type": "object"