⚠️ fix(datashare-python): make Document.root_document not nullable and Document.extraction_level mandatory

datashare-python/datashare_python/conftest.py

@@ -233,14 +233,22 @@ def text_1() -> str:
 @pytest.fixture(scope="session")
 def doc_0(text_0: str) -> Document:
     return Document(
-        id="doc-0", root_document="root-0", language="ENGLISH", content=text_0
+        id="doc-0",
+        root_document="root-0",
+        extraction_level=1,
+        language="ENGLISH",
+        content=text_0,
     )
 
 
 @pytest.fixture(scope="session")
 def doc_1(text_1: str) -> Document:
     return Document(
-        id="doc-1", root_document="root-1", language="ENGLISH", content=text_1
+        id="doc-1",
+        root_document="root-1",
+        extraction_level=1,
+        language="ENGLISH",
+        content=text_1,
     )
 
 
@@ -249,6 +257,7 @@ def doc_2() -> Document:
     return Document(
         id="doc-2",
         root_document="root-2",
+        extraction_level=1,
         language="FRENCH",
         content="traduis ce texte en anglais",
     )
@@ -260,6 +269,7 @@ def doc_3() -> Document:
         id="doc-3",
         index=TEST_PROJECT,
         root_document="root-3",
+        extraction_level=1,
         language="SPANISH",
         content="traduce este texto al inglés",
     )

datashare-python/datashare_python/objects.py

@@ -24,6 +24,7 @@
     from icij_common.es import (
         DOC_CONTENT,
         DOC_CONTENT_TRANSLATED,
+        DOC_EXTRACTION_LEVEL,
         DOC_LANGUAGE,
         DOC_METADATA,
         DOC_PATH,
@@ -222,6 +223,8 @@ class Document(DatashareModel):
     language: DatashareLanguage
     index: str | None = None
     root_document: str | None = None
+    content: Annotated[str, BeforeValidator(_from_sentences)]
+    extraction_level: int = 0
     content: str | None = None
     content_text_length: int | None = None
     content_type: str | None = None
@@ -244,6 +247,37 @@ def _initialize_content_length_from_content(cls, data: Any) -> Any:
                 data["content_text_length"] = len(content)
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def _set_root_document_to_self_when_not_provided(cls, data: Any) -> Any:
+        if isinstance(data, dict):
+            root = data.get("root_document")
+            if root is None:
+                data["root_document"] = data.get("id")
+        return data
+
+    @model_validator(mode="after")
+    def _validate_root_document(self) -> Any:
+        if self.root_document is None:
+            msg = (
+                "inconsistent state, root_document cannot be after validation, "
+                "there's a bug in this object's implementation."
+                " Please report to a developer."
+            )
+            raise ValueError(msg)
+        if self.is_root_document:
+            if self.extraction_level > 0:
+                msg = "extraction_level should be <= 0 for root documents"
+                raise ValueError(msg)
+        elif self.extraction_level <= 0:
+            msg = "extraction_level should be > 0 for embedded documents"
+            raise ValueError(msg)
+        return self
+
+    @property
+    def is_root_document(self) -> bool:
+        return self.root_document == self.id and self.extraction_level <= 0
+
     @classmethod
     def from_es(cls, es_doc: dict) -> Self:
         sources = es_doc[SOURCE]
@@ -255,6 +289,7 @@ def from_es(cls, es_doc: dict) -> Self:
             content_text_length=sources.get("content_text_length"),
             language=DatashareLanguage(sources[DOC_LANGUAGE]),
             root_document=sources.get(DOC_ROOT_ID),
+            extraction_level=sources.get(DOC_EXTRACTION_LEVEL),
             tags=sources.get("tags", []),
             path=sources.get(DOC_PATH),
             metadata=sources.get(DOC_METADATA),
@@ -275,7 +310,7 @@ def to_processed_file(self) -> ProcessedFile:
             )
             raise ValueError(msg)
         resource_name = cast(str, resource_name)
-        if self.root_document is None:
+        if self.is_root_document:
             path = self.path
             location = DocumentLocation.FILESYSTEM
         else:

datashare-python/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datashare-python"
-version = "0.10.10"
+version = "0.11.0"
 description = "Manage Python tasks and local resources in Datashare"
 authors = [
   { name = "Clément Doumouro", email = "cdoumouro@icij.org" },

datashare-python/tests/test_objects.py

@@ -57,6 +57,58 @@ def test_task_ser() -> None:
     assert serialized == expected
 
 
+def test_embedded_document_should_have_extraction_level() -> None:
+    # When/Then
+    expected = "extraction_level should be > 0 for embedded documents"
+    with pytest.raises(ValidationError, match=expected):
+        Document(
+            id="some_id", root_document="root_id", language=DatashareLanguage("ENGLISH")
+        )
+
+
+def test_embedded_document_should_have_extraction_level_greater_than_0() -> None:
+    # When/Then
+    expected = "extraction_level should be > 0 for embedded documents"
+    with pytest.raises(ValidationError, match=expected):
+        Document(
+            id="some_id",
+            root_document="root_id",
+            language=DatashareLanguage("ENGLISH"),
+            extraction_level=0,
+        )
+
+
+def test_embedded_document_should_have_consistent_extraction_level() -> None:
+    # When/Then
+    expected = "extraction_level should be > 0 for embedded documents"
+    with pytest.raises(ValidationError, match=expected):
+        Document(
+            id="some_id",
+            root_document="root_id",
+            language=DatashareLanguage("ENGLISH"),
+            extraction_level=0,
+        )
+
+
+@pytest.mark.parametrize(
+    ("doc", "is_root"),
+    [
+        (Document(id="some_id", language=DatashareLanguage("ENGLISH")), True),
+        (
+            Document(
+                id="some_id",
+                root_document="root_id",
+                extraction_level=1,
+                language=DatashareLanguage("ENGLISH"),
+            ),
+            False,
+        ),
+    ],
+)
+def test_is_root_document(doc: Document, *, is_root: bool) -> None:
+    assert doc.is_root_document == is_root
+
+
 def test_filesystem_document_should_raise_on_absolute_path() -> None:
     # Given
     path = Path("/some/absolute/path")
@@ -78,7 +130,11 @@ def test_document_to_filesystem_document_use_relative_path() -> None:
     assert path.is_absolute()
     meta = {TIKA_METADATA_RESOURCENAME: "resource.file"}
     doc = Document(
-        index=TEST_PROJECT, path=path, id="some_id", language="ENGLISH", metadata=meta
+        index=TEST_PROJECT,
+        path=path,
+        id="some_id",
+        language=DatashareLanguage("ENGLISH"),
+        metadata=meta,
     )
     # When
     fs_doc = doc.to_processed_file()

worker-template/.python-version

@@ -1 +1 @@
-3.11
+3.12

worker-template/pyproject.toml

@@ -7,7 +7,7 @@ authors = [
   { name = "Clément Doumouro", email = "clement.doumouro@gmail.com" }
 ]
 readme = "README.md"
-requires-python = ">=3.11,<4"
+requires-python = ">=3.12,<4"
 dependencies = [
   "datashare-python~=0.10.0",
   "pycountry~=26.2.16",