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

datashare_python/objects.py

@@ -1,22 +1,22 @@
 import logging
 import os
-from asyncio import Lock
+from abc import ABC
 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
 from pathlib import Path
-from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, cast
+from typing import Annotated, Any, ClassVar, Generic, Literal, Self, TypeVar, cast
 
 import langcodes
-from lru import LRU
+from icij_common.registrable import Registrable
 from pydantic_core import PydanticCustomError, ValidationError, core_schema
 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 (
@@ -34,12 +34,17 @@ with workflow.unsafe.imports_passed_through():
 from icij_common.pydantic_utils import (
     icij_config,
     lowercamel_case_config,
+    make_enum_discriminator,
     merge_configs,
     no_enum_values_config,
+    tagged_union,
 )
 from pydantic import (
     AfterValidator,
+    AliasChoices,
     BeforeValidator,
+    ConfigDict,
+    Discriminator,
     Field,
     GetCoreSchemaHandler,
     TypeAdapter,
@@ -104,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"
 
@@ -119,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__(
@@ -174,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
@@ -217,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 = (
@@ -241,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():
@@ -256,13 +300,142 @@ def _is_absolute_path(v: bytes | BytesIO | Path) -> Any:
     return v
 
 
-@dataclass(frozen=True)
-class DocArtifact:
+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
+        # Dump in json mode to make testing easier
+        as_manifest = self.model_dump(by_alias=True, mode="json")
+        return as_manifest
+
+
+A = TypeVar("A", bound=TaskArgs)
+
+
+class ManifestEntry(DatashareModel, Generic[A], ABC):
+    status: ManifestEntryStatus
+    label: str | None = None
+    input: Annotated[
+        dict[str, Any] | None,
+        Field(
+            validation_alias=AliasChoices("taskInput", "input"),
+            serialization_alias="taskInput",
+        ),
+    ]
+
+    @classmethod
+    def complete(cls, args: A, label: str | None = None, **kwargs) -> Self:
+        return cls(
+            input=args.as_manifest_task_input(),
+            label=label,
+            status=ManifestEntryStatus.COMPLETE,
+            **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"
+
+
+def _validate_pages_range(v: Any) -> None:
+    if not isinstance(v, list):
+        msg = f"expected a list, got {type(v)}"
+        raise TypeError(msg)
+    previous_end = None
+    for page_i, (start, end) in enumerate(v):
+        if not start <= end:
+            msg = "end of page must be >= start"
+            raise ValueError(msg)
+        if previous_end is not None and previous_end != start:
+            msg = (
+                f"start of page {page_i} doesn't match end of previous "
+                f"page {previous_end}"
+            )
+            raise ValueError(msg)
+    return v
+
+
+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(
+        default=PaginationType.FILESYSTEM, frozen=True
+    )
+
+
+@BasePagination.register(PaginationType.BYTE_RANGES)
+class ByteRangesPagination(BasePagination):
+    type: ClassVar[PaginationType] = Field(
+        default=PaginationType.BYTE_RANGES, frozen=True
+    )
+    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 (
+            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 {n_pages} for"
+                f" byte_ranges and  {self.total} for total."
+            )
+            raise ValueError(msg)
+        return self
+
+
+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)
+
     project: str
     doc_id: str
     artifact: Annotated[bytes | BytesIO | Path, AfterValidator(_is_absolute_path)]
-    filename: str
-    metadata_key: str
+    filename: ClassVar[str]  # Override this
+    type: ClassVar[ArtifactType]  # Override this
+    manifest_entry: ManifestEntry
 
 
 @unique
@@ -369,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)