datashare-python 0.9.9__py3-none-any.whl → 0.10.0.dev0__py3-none-any.whl

datashare_python/objects.py

@@ -1,9 +1,8 @@
 import logging
 import os
 from abc import ABC
-from asyncio import Lock
 from collections.abc import Awaitable, Callable
-from dataclasses import InitVar, dataclass, field
+from dataclasses import dataclass
 from datetime import UTC, datetime
 from enum import StrEnum, unique
 from io import BytesIO
@@ -12,15 +11,12 @@ from typing import Annotated, Any, ClassVar, Generic, Literal, Self, TypeVar, ca
 
 import langcodes
 from icij_common.registrable import Registrable
-from lru import LRU
 from pydantic_core import PydanticCustomError, ValidationError, core_schema
-from pydantic_core.core_schema import (
-    PlainValidatorFunctionSchema,
-)
+from pydantic_core.core_schema import PlainValidatorFunctionSchema
 from pydantic_extra_types.language_code import LanguageName
 from temporalio import workflow
 
-from .constants import DEFAULT_SHARED_RESOURCES_SIZE, TIKA_METADATA_RESOURCENAME
+from .constants import TIKA_METADATA_RESOURCENAME
 
 with workflow.unsafe.imports_passed_through():
     from icij_common.es import (
@@ -113,9 +109,15 @@ class DatashareLanguage(str):
         return self.as_language_name.alpha3
 
 
+class WorkerPaths(BaseModel):
+    filesystem: Path
+    artifacts: Path
+    workdir: Path
+
+
 @unique
 class DocumentLocation(StrEnum):
-    ORIGINAL = "original"
+    FILESYSTEM = "filesystem"
     ARTIFACTS = "artifacts"
     WORKDIR = "workdir"
 
@@ -128,30 +130,49 @@ def _is_relative(value: Path) -> Path:
     return value
 
 
-class FilesystemDocument(DatashareModel):
+class ProcessedFile(BaseModel):
     id: str
     path: Annotated[Path, AfterValidator(_is_relative)]
-    index: str
+    project: str
     location: DocumentLocation
     resource_name: str
+    n_pages: int
+    parent: "ProcessedFile | None" = None
+
+    @classmethod
+    def from_doc(cls, doc: "Document") -> Self:
+        return doc.to_processed_file()
+
+    def child(self, path: Path, paths: WorkerPaths) -> Self:
+        return ProcessedFile(
+            id=self.id,
+            path=path.relative_to(paths.workdir),
+            project=self.project,
+            location=DocumentLocation.WORKDIR,
+            resource_name=path.name,
+            n_pages=self.n_pages,
+            parent=self,
+        )
 
-    def locate(
-        self, original_root: Path, *, artifacts_root: Path, workdir: Path
-    ) -> Path:
+    def locate(self, paths: WorkerPaths) -> Path:
         from datashare_python.utils import artifacts_dir  # noqa: PLC0415
 
         match self.location:
-            case DocumentLocation.ORIGINAL:
-                return original_root / self.path
+            case DocumentLocation.FILESYSTEM:
+                return paths.filesystem / self.path
             case DocumentLocation.ARTIFACTS:
-                project = self.index
-                return artifacts_root / artifacts_dir(self.id, project=project) / "raw"
+                project = self.project
+                return paths.artifacts / artifacts_dir(self.id, project=project) / "raw"
             case DocumentLocation.WORKDIR:
-                return workdir / self.path
+                return paths.workdir / self.path
             case _:
                 raise ValueError(f"invalid location: {self.path}")
 
 
+class ProcessedPage(ProcessedFile):
+    page_number: int
+
+
 class IETFLanguage(str):
     @classmethod
     def __get_pydantic_core_schema__(
@@ -183,6 +204,11 @@ class Translation(BaseModel):  # No camelcase here we don't know why
     content: Annotated[str, BeforeValidator(_from_sentences)]
 
 
+Routing = str
+DocID = str
+DocRoute = tuple[DocID, Routing]
+
+
 class Document(DatashareModel):
     id: str
     language: DatashareLanguage
@@ -226,17 +252,17 @@ class Document(DatashareModel):
             metadata=sources.get(DOC_METADATA),
         )
 
-    def to_filesystem(self) -> FilesystemDocument:
+    def to_processed_file(self) -> ProcessedFile:
         from .utils import artifacts_dir  # noqa: PLC0415
 
         if self.metadata is None:
             raise ValueError(
-                "can't compute filesyste path for document withtout metadata"
+                "can't compute filesystem path for document withtout metadata"
             )
         resource_name = cast(str, self.metadata[TIKA_METADATA_RESOURCENAME])
         if self.root_document is None:
             path = self.path
-            location = DocumentLocation.ORIGINAL
+            location = DocumentLocation.FILESYSTEM
         else:
             if self.index is None:
                 msg = (
@@ -250,14 +276,23 @@ class Document(DatashareModel):
         # we store a relative path otherwise joining with the location will fail
         if path.parts and path.parts[0] == os.path.sep:
             path = Path(*path.parts[1:])
-        return FilesystemDocument(
+        n_pages = 1
+        if self.metadata:
+            n_pages = self.metadata.get("tika_metadata_xmptpg_npages", n_pages)
+        return ProcessedFile(
             id=self.id,
             path=path,
-            index=self.index,
+            project=self.index,
             location=location,
             resource_name=resource_name,
+            n_pages=n_pages,
         )
 
+    def to_route(self) -> DocRoute:
+        if self.root_document is not None:
+            return self.root_document, self.id
+        return self.id, self.id
+
 
 def _is_absolute_path(v: bytes | BytesIO | Path) -> Any:
     if isinstance(v, Path) and not v.is_absolute():
@@ -268,17 +303,20 @@ def _is_absolute_path(v: bytes | BytesIO | Path) -> Any:
 class ArtifactType(StrEnum):
     STRUCTURE = "structure"
     ASR_TRANSCRIPTION = "transcription"
+    PASSPORTS = "passports"
 
 
 class ManifestEntryStatus(StrEnum):
     COMPLETE = "complete"
+    PARTIAL = "partial"
 
 
 class TaskArgs(DatashareModel, ABC):
     def as_manifest_task_input(self) -> dict[str, Any]:
         # This is a base implementation, if the input is too large to be dumped,
         # override this and pop large keys
-        as_manifest = self.model_dump(by_alias=True)
+        # Dump in json mode to make testing easier
+        as_manifest = self.model_dump(by_alias=True, mode="json")
         return as_manifest
 
 
@@ -305,19 +343,21 @@ class ManifestEntry(DatashareModel, Generic[A], ABC):
             **kwargs,
         )
 
+    @classmethod
+    def partial(cls, args: A, label: str | None = None, **kwargs) -> Self:
+        return cls(
+            input=args.as_manifest_task_input(),
+            label=label,
+            status=ManifestEntryStatus.PARTIAL,
+            **kwargs,
+        )
+
 
 class PaginationType(StrEnum):
     FILESYSTEM = "filesystem"
     BYTE_RANGES = "byteRanges"
 
 
-class BasePagination(DatashareModel, Registrable, ABC):
-    registry_key: ClassVar[str] = Field(frozen=True, default="type")
-
-    total: int
-    type: ClassVar[PaginationType] = Field(frozen=True)
-
-
 def _validate_pages_range(v: Any) -> None:
     if not isinstance(v, list):
         msg = f"expected a list, got {type(v)}"
@@ -339,6 +379,11 @@ def _validate_pages_range(v: Any) -> None:
 PagesRange = Annotated[list[tuple[int, int]], AfterValidator(_validate_pages_range)]
 
 
+class BasePagination(Registrable, DatashareModel):
+    registry_key: ClassVar[str] = Field(frozen=True, default="type")
+    type: ClassVar[PaginationType] = Field(frozen=True)
+
+
 @BasePagination.register(PaginationType.FILESYSTEM)
 class FilesystemPagination(BasePagination):
     type: ClassVar[PaginationType] = Field(
@@ -353,25 +398,34 @@ class ByteRangesPagination(BasePagination):
     )
     byte_ranges: PagesRange
 
+
+pagination_discriminator = make_enum_discriminator("type", PaginationType)
+Pagination = Annotated[
+    tagged_union(BasePagination.__subclasses__(), lambda x: x.type.default),
+    Discriminator(pagination_discriminator),
+]
+
+
+class Pages(DatashareModel):
+    total: int
+    pagination: Pagination
+
     @model_validator(mode="after")
     def byte_ranges_length_should_match_total(self) -> Self:
-        if len(self.byte_ranges) != self.total:
+        if (
+            isinstance(self.pagination, ByteRangesPagination)
+            and len(self.pagination.byte_ranges) != self.total
+        ):
+            n_pages = len(self.pagination.byte_ranges)
             msg = (
-                f"byte_ranges must match total. Found {len(self.byte_ranges)} for"
+                f"byte_ranges must match total. Found {n_pages} for"
                 f" byte_ranges and  {self.total} for total."
             )
             raise ValueError(msg)
         return self
 
 
-pagination_discriminator = make_enum_discriminator("type", PaginationType)
-Pagination = Annotated[
-    tagged_union(BasePagination.__subclasses__(), lambda x: x.type.default),
-    Discriminator(pagination_discriminator),
-]
-
-
-class DocArtifact(BaseModel, ABC):
+class DocArtifact(DatashareModel, ABC):
     # This object is not used for serde, just as a container, it's OK to allow
     # arbitrary types (to allow storing BytesIO)
     model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -488,47 +542,3 @@ class TaskGroup:
     @classmethod
     def python(cls) -> Self:
         return cls(name="PYTHON")
-
-
-@dataclass(frozen=True)
-class Shared:
-    cache_size: InitVar[int] = DEFAULT_SHARED_RESOURCES_SIZE
-    eviction_callback: InitVar[Callable] = None
-    _resources: LRU = field(init=False, repr=False)
-    _lock: Lock = field(init=False, repr=False)
-
-    def __post_init__(self, cache_size: int, eviction_callback: Callable) -> None:
-        object.__setattr__(
-            self, "_resources", LRU(cache_size, callback=eviction_callback)
-        )
-        object.__setattr__(self, "_lock", Lock())
-
-    def get_resource(
-        self, key: str, default: Any = None, *, set_if_unavailable: bool = True
-    ) -> Any:
-        if key not in self._resources and set_if_unavailable:
-            self.set_resource(key, default)
-
-        return self._resources.get(key, default)
-
-    def set_resource(self, key: str, value: Any) -> None:
-        self._resources[key] = value
-
-    def pop_resource(self, key: str, default: Any = None) -> Any:
-        return self._resources.pop(key, default)
-
-    async def async_get_resource(
-        self, key: str, default: Any = None, *, set_if_unavailable: bool = True
-    ) -> Any:
-        if key not in self._resources and set_if_unavailable:
-            await self.async_set_resource(key, default)
-
-        return self._resources.get(key, default)
-
-    async def async_set_resource(self, key: str, value: Any) -> None:
-        async with self._lock:
-            self._resources[key] = value
-
-    async def async_pop_resource(self, key: str, default: Any = None) -> Any:
-        async with self._lock:
-            return self._resources.pop(key, default)