Simpler literal enforcement

Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
DS4SD · Sep 30, 2024 · 46a6d8d · 46a6d8d
1 parent 677c759
commit 46a6d8d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 13 deletions.
diff --git a/docling_core/types/experimental/document.py b/docling_core/types/experimental/document.py
@@ -7,7 +7,6 @@
 
 import pandas as pd
 from pydantic import (
-    AfterValidator,
     AnyUrl,
     BaseModel,
     ConfigDict,
@@ -324,9 +323,7 @@ def export_to_document_tokens(
 class SectionHeaderItem(TextItem):
     """SectionItem."""
 
-    label: typing.Annotated[
-        DocItemLabel, AfterValidator(lambda x: DocItemLabel.SECTION_HEADER)
-    ] = Field(default=DocItemLabel.SECTION_HEADER, frozen=True)
+    label: typing.Literal[DocItemLabel.SECTION_HEADER] = DocItemLabel.SECTION_HEADER
     level: LevelNumber
 
 
@@ -342,9 +339,8 @@ class FloatingItem(DocItem):
 class PictureItem(FloatingItem):
     """PictureItem."""
 
-    label: typing.Annotated[
-        DocItemLabel, AfterValidator(lambda x: DocItemLabel.PICTURE)
-    ] = Field(default=DocItemLabel.PICTURE, frozen=True)
+    label: typing.Literal[DocItemLabel.PICTURE] = DocItemLabel.PICTURE
+
     data: BasePictureData
 
     def export_to_document_tokens(
@@ -401,9 +397,7 @@ class TableItem(FloatingItem):
     """TableItem."""
 
     data: BaseTableData
-    label: typing.Annotated[
-        DocItemLabel, AfterValidator(lambda x: DocItemLabel.TABLE)
-    ] = Field(default=DocItemLabel.TABLE, frozen=True)
+    label: typing.Literal[DocItemLabel.TABLE] = DocItemLabel.TABLE
 
     def export_to_dataframe(self) -> pd.DataFrame:
         """Export the table as a Pandas DataFrame."""

diff --git a/test/test_docling_doc.py b/test/test_docling_doc.py
@@ -89,23 +89,21 @@ def verify(dc, obj):
             obj = dc(
                 text="whatever",
                 orig="whatever",
-                label=DocItemLabel.TEXT,
+                label=DocItemLabel.SECTION_HEADER,
                 self_ref="#",
                 level=2,
             )
             verify(dc, obj)
 
         elif dc is PictureItem:
             obj = dc(
-                label=DocItemLabel.TEXT,
                 self_ref="#",
                 data=BasePictureData(),
             )
             verify(dc, obj)
 
         elif dc is TableItem:
             obj = dc(
-                label=DocItemLabel.TEXT,
                 self_ref="#",
                 data=BaseTableData(num_rows=3, num_cols=5, table_cells=[]),
             )