datashare-python 0.11.3__tar.gz → 0.12.0__tar.gz

{datashare_python-0.11.3 → datashare_python-0.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: datashare-python
-Version: 0.11.3
+Version: 0.12.0
 Summary: Manage Python tasks and local resources in Datashare
 Project-URL: Homepage, https://icij.github.io/datashare-python/
 Project-URL: Documentation, https://icij.github.io/datashare-python/

{datashare_python-0.11.3 → datashare_python-0.12.0}/datashare_python/mimetypes_.py

@@ -75,6 +75,7 @@ def types_map() -> dict[str, set[str]]:
         ".bmp": {"image/bmp"},
         ".book": {"application/vnd.framemaker"},
         ".box": {"application/vnd.previewsystems.box"},
+        ".boxnote": {"application/octet-stream"},
         ".boz": {"application/x-bzip2"},
         ".bpk": {"application/octet-stream"},
         ".bufr": {"application/x-bufr"},
@@ -172,6 +173,7 @@ def types_map() -> dict[str, set[str]]:
         ".davmount": {"application/davmount+xml"},
         ".dbf": {"application/dbf", "application/dbase"},
         ".dbk": {"application/docbook+xml"},
+        ".dclx": {"application/octet-stream"},
         ".dcr": {"application/x-director"},
         ".dcurl": {"text/vnd.curl.dcurl"},
         ".dcx": {"image/vnd.zbrush.dcx", "image/x-dcx"},
@@ -226,6 +228,8 @@ def types_map() -> dict[str, set[str]]:
         ".dxf": {"image/vnd.dxf"},
         ".dxp": {"application/vnd.spotfire.dxp"},
         ".dxr": {"application/x-director"},
+        ".ebc": {"application/octet-stream"},
+        ".ebcdic": {"application/octet-stream"},
         ".ecelp4800": {"audio/vnd.nuera.ecelp4800"},
         ".ecelp7470": {"audio/vnd.nuera.ecelp7470"},
         ".ecelp9600": {"audio/vnd.nuera.ecelp9600"},

{datashare_python-0.11.3 → datashare_python-0.12.0}/datashare_python/objects.py

@@ -2,14 +2,15 @@ import hashlib
 import json
 import logging
 import os
-from abc import ABC
-from collections.abc import Awaitable, Callable
+import traceback
+from abc import ABC, abstractmethod
+from collections.abc import Callable
 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, Literal, Self, TypeVar, cast, final
 
 import langcodes
 from icij_common.registrable import Registrable
@@ -50,6 +51,7 @@ from pydantic import (
     Discriminator,
     Field,
     GetCoreSchemaHandler,
+    Tag,
     TypeAdapter,
     model_validator,
 )
@@ -59,8 +61,9 @@ from pydantic.main import IncEx
 logger = logging.getLogger(__name__)
 
 
-T = TypeVar("T")
-Predicate = Callable[[T], bool] | Callable[[T], Awaitable[bool]]
+Routing = str
+DocID = str
+DocRoute = tuple[DocID, Routing]
 
 
 class BaseModel(_BaseModel):
@@ -118,68 +121,34 @@ class DatashareLanguage(str):
         return self.as_language_name.alpha3
 
 
-class WorkerPaths(BaseModel):
-    filesystem: Path
-    artifacts: Path
-    workdir: Path
-
-
 @unique
-class DocumentLocation(StrEnum):
+class FileLocation(StrEnum):
     FILESYSTEM = "filesystem"
     ARTIFACTS = "artifacts"
     WORKDIR = "workdir"
 
 
-def _is_relative(value: Path) -> Path:
-    if value.is_absolute():
-        raise ValueError(
-            f"FilesystemDocument path should always be relative, found {value}"
-        )
-    return value
-
-
-class ProcessedFile(BaseModel):
-    id: str
-    path: Annotated[Path, AfterValidator(_is_relative)]
-    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,
-        )
+class WorkerPaths(BaseModel):
+    filesystem: Path
+    artifacts: Path
+    workdir: Path
 
-    def locate(self, paths: WorkerPaths) -> Path:
-        from datashare_python.utils import artifacts_dir  # noqa: PLC0415
-
-        match self.location:
-            case DocumentLocation.FILESYSTEM:
-                return paths.filesystem / self.path
-            case DocumentLocation.ARTIFACTS:
-                project = self.project
-                return paths.artifacts / artifacts_dir(self.id, project=project) / "raw"
-            case DocumentLocation.WORKDIR:
-                return paths.workdir / self.path
+    def locate(self, path: Path, location: FileLocation) -> Path:
+        match location:
+            case FileLocation.FILESYSTEM:
+                return self.filesystem / path
+            case FileLocation.ARTIFACTS:
+                return self.artifacts / path
+            case FileLocation.WORKDIR:
+                return self.workdir / path
             case _:
-                raise ValueError(f"invalid location: {self.path}")
+                raise ValueError(f"invalid location: {path}")
 
 
-class ProcessedPage(ProcessedFile):
-    page_number: int
+class FromParent[P](ABC):
+    @classmethod
+    @abstractmethod
+    def from_parent(cls, parent: P, *args, **kwargs) -> Self: ...
 
 
 class IETFLanguage(str):
@@ -213,12 +182,9 @@ 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):
+    """ES datashare document"""
+
     id: str
     language: DatashareLanguage
     index: str | None = None
@@ -285,7 +251,7 @@ class Document(DatashareModel):
             id=es_doc[ID_],
             index=es_doc.get(INDEX_),
             content=sources.get(DOC_CONTENT),
-            content_translated=sources.get(DOC_CONTENT_TRANSLATED, []),
+            content_translated=sources.get(DOC_CONTENT_TRANSLATED),
             content_text_length=sources.get("content_text_length"),
             language=DatashareLanguage(sources[DOC_LANGUAGE]),
             root_document=sources.get(DOC_ROOT_ID),
@@ -295,58 +261,57 @@ class Document(DatashareModel):
             metadata=sources.get(DOC_METADATA),
         )
 
-    def to_processed_file(self) -> ProcessedFile:
-        from .utils import artifacts_dir  # noqa: PLC0415
+    @property
+    def project(self) -> str:
+        if self.index is None:
+            raise ValueError("missing index")
+        return self.index
 
+    @property
+    def doc_id(self) -> str:
+        return self.id
+
+    @property
+    def route(self) -> DocRoute:
+        if self.root_document:
+            return self.id, self.id
+        return self.id, self.root_document
+
+    @property
+    def resource_name(self) -> str:
         if self.metadata is None:
-            raise ValueError(
-                "can't compute filesystem path for document withtout metadata"
-            )
+            raise ValueError("missing metadata")
         resource_name = self.metadata.get(TIKA_METADATA_RESOURCENAME)
         if resource_name is None:
-            # Backward compat for project before resourcename was mandatory
-            if self.is_root_document:
-                resource_name = self.path.name
-            else:
-                msg = (
-                    f"doc {self} can't turn ES document into a"
-                    f" {ProcessedFile.__name__} without tika metadata resources name"
-                    f" in metadata: {self.metadata}"
-                )
-                raise ValueError(msg)
-        resource_name = cast(str, resource_name)
-        if self.is_root_document:
-            path = self.path
-            location = DocumentLocation.FILESYSTEM
-        else:
-            if self.index is None:
-                msg = (
-                    f"can't compute filesystem path for embedded doc {self.id} without"
-                    f" index"
-                )
-                raise ValueError(msg)
-            path = artifacts_dir(doc_id=self.id, project=self.index) / "raw"
-            location = DocumentLocation.ARTIFACTS
-        # The filesystem dod is alway relative to the base location, let's make sure
-        # 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:])
-        n_pages = 1
-        if self.metadata:
-            n_pages = self.metadata.get("tika_metadata_xmptpg_npages", n_pages)
-        return ProcessedFile(
-            id=self.id,
-            path=path,
-            project=self.index,
-            location=location,
-            resource_name=resource_name,
-            n_pages=n_pages,
+            msg = f"missing {TIKA_METADATA_RESOURCENAME} in metadata"
+            raise KeyError(msg)
+        cast(str, resource_name)
+        return resource_name
+
+
+def _is_relative(value: Path) -> Path:
+    if value.is_absolute():
+        raise ValueError(
+            f"WorkerFilePath path should always be relative, found {value}"
         )
+    return value
 
-    def to_route(self) -> DocRoute:
-        if self.root_document is not None:
-            return self.root_document, self.id
-        return self.id, self.id
+
+class WorkerFilePath(BaseModel):
+    """Path of a file produced or processed by a worker, located somewhere on the DS
+    filesystem, artifactsdir or workdir
+    """
+
+    path: Annotated[Path, AfterValidator(_is_relative)]
+    location: FileLocation
+
+    def locate(self, paths: WorkerPaths) -> Path:
+        return paths.locate(self.path, self.location)
+
+    @classmethod
+    def relative_to_workdir(cls, path: Path, paths: WorkerPaths) -> Self:
+        path = path.relative_to(paths.workdir)
+        return cls(path=path, location=FileLocation.WORKDIR)
 
 
 def _is_absolute_path(v: bytes | BytesIO | Path) -> Any:
@@ -617,3 +582,297 @@ class TaskGroup:
     @classmethod
     def python(cls) -> Self:
         return cls(name="PYTHON")
+
+
+class Error(BaseModel):
+    title: str
+    detail: str | None
+
+    @classmethod
+    def from_exception(cls, exception: BaseException) -> "Error":
+        title = exception.__class__.__name__
+        trace_lines = traceback.format_exception(
+            None, value=exception, tb=exception.__traceback__
+        )
+        detail = f"{exception}\n{''.join(trace_lines)}"
+        error = Error(title=title, detail=detail)
+        return error
+
+    def without_detail(self) -> "Error":
+        return Error(title=self.title, detail=None)
+
+
+class ErrorSource(BaseModel, ABC):
+    """Source of the processing error."""
+
+    doc_id: str
+    project: str
+
+
+class ErrorSourceWithPages(ErrorSource, ABC):
+    """Source of the processing error and number of pages processed at that time"""
+
+    n_pages: int
+
+
+class ProcessingError[S: ErrorSource](BaseModel):
+    """Error occurring while processing a document or intermediate result"""
+
+    source: S
+    error: Error
+
+    @final
+    @classmethod
+    def from_exception(cls, source: S, exception: BaseException) -> Self:
+        return cls(source=source, error=Error.from_exception(exception))
+
+
+class ProcessingReport(DatashareModel):
+    """Basic processing report indicating how many docs where processed"""
+
+    n_docs: int = 0
+
+    def __add__(self, other: Self) -> Self:
+        return ProcessingReport(n_docs=other.n_docs + self.n_docs)
+
+
+class ProcessingReportWithPages(ProcessingReport):
+    """Processing report indicating how many pages where processed"""
+
+    n_pages: int = 0
+
+    def __add__(self, other: Self) -> Self:
+        return ProcessingReportWithPages(
+            n_docs=other.n_docs + self.n_docs, n_pages=other.n_pages + self.n_pages
+        )
+
+
+def _without_detail(errors: list[Error]) -> list[Error]:
+    if any(e.detail is not None for e in errors):
+        msg = "expected errors without details"
+        raise ValueError(msg)
+    return errors
+
+
+class DocProcessingErrors(BaseModel):
+    """All errors related to several processing happening on a single document.
+    If a source doc has many children (typically a document split into pages
+    processed separately), all children errors are merged under the older parent.
+    """
+
+    doc_id: str
+    root_document: str
+    project: str
+    errors: Annotated[
+        list[Error], AfterValidator(_without_detail), Field(default_factory=list)
+    ]
+
+
+class ErrorReport(ProcessingReport):
+    """Processing report indicated how many docs failed and listing errors per doc."""
+
+    errors: list[DocProcessingErrors] = Field(default_factory=list)
+
+    def __add__(self, other: Self) -> Self:
+        return ErrorReportWithPages(
+            n_docs=other.n_docs + self.n_docs, errors=self.errors + other.errors
+        )
+
+    @classmethod
+    def from_errors(cls, *errors: ProcessingError) -> Self:
+        """Aggregate errors by root source documents"""
+        roots = dict()
+        for error in errors:
+            root = error.source
+            while (parent := getattr(root, "parent", None)) is not None:
+                root = parent
+            root_errors = roots.get(root.doc_id)
+            if root_errors is not None:
+                _, root_errors = root_errors
+            else:
+                root_errors = []
+            root_errors.append(error)
+            roots[root.doc_id] = (root, root_errors)
+        errors = []
+        for _, (root, root_errors) in sorted(roots.items()):
+            # We keep details for errors stored as intermediate steps to allow
+            # debugging, we don't output them in the response to keep everything
+            # lightweight and more secure
+            root_errors = [err.error.without_detail() for err in root_errors]  # noqa: PLW2901
+            doc_errors = DocProcessingErrors(
+                doc_id=root.doc_id,
+                project=root.project,
+                root_document=root.root_document,
+                errors=root_errors,
+            )
+            errors.append(doc_errors)
+        return cls(n_docs=len(errors), errors=errors)
+
+
+class ErrorReportWithPages(ProcessingReportWithPages, ABC):
+    """Same as ErrorReport but providing details about how many pages failed"""
+
+    errors: list[DocProcessingErrors] = Field(default_factory=list)
+
+    def __add__(self, other: Self) -> Self:
+        return ErrorReportWithPages(
+            n_docs=other.n_docs + self.n_docs,
+            n_pages=other.n_pages + self.n_pages,
+            errors=self.errors + other.errors,
+        )
+
+    @classmethod
+    def from_errors(cls, *errors: ProcessingError) -> Self:
+        """Aggregate errors by root source documents"""
+        roots = dict()
+        n_pages = 0
+        for error in errors:
+            root = error.source
+            while (parent := getattr(root, "parent", None)) is not None:
+                root = parent
+            root_errors = roots.get(root.doc_id)
+            if root_errors is not None:
+                _, root_errors = root_errors
+            else:
+                root_errors = []
+            n_pages += error.source.n_pages
+            root_errors.append(error)
+            roots[root.doc_id] = (root, root_errors)
+        errors = []
+        for _, (root, root_errors) in sorted(roots.items()):
+            # We keep details for errors stored as intermediate steps to allow
+            # debugging, we don't output them in the response to keep everything
+            # lightweight and more secure
+            root_errors = [err.error.without_detail() for err in root_errors]  # noqa: PLW2901
+            doc_errors = DocProcessingErrors(
+                doc_id=root.doc_id,
+                project=root.project,
+                root_document=root.root_document,
+                errors=root_errors,
+            )
+            errors.append(doc_errors)
+        return cls(n_docs=len(errors), n_pages=n_pages, errors=errors)
+
+
+class ProcessingResult[V](BaseModel):
+    value: V
+
+
+class WithParent[P](FromParent[P]):
+    """Result produced from another processing result
+
+    Tracking parent allows track errors from one intermediate results to another back
+     to the initial DS Document from the index. It also allows error aggregation.
+    """
+
+    parent: P
+
+    @property
+    def doc_id(self) -> str:
+        return self.parent.doc_id
+
+    @property
+    def root_document(self) -> str:
+        return self.parent.root_document
+
+    @property
+    def project(self) -> str:
+        return self.parent.project
+
+    @property
+    def route(self) -> DocRoute:
+        return self.parent.route
+
+
+class DatashareFile(
+    WithParent[Document], ProcessingResult[WorkerFilePath], ErrorSourceWithPages
+):
+    """A datashare document file ready to be processed and located somewhere on the
+    filesystem or artifact directory.
+    """
+
+    parent: Document
+
+    def locate(self, paths: WorkerPaths) -> Path:
+        return self.value.locate(paths)
+
+    @classmethod
+    def from_parent(cls, parent: Document) -> Self:
+        from .utils import artifacts_dir  # noqa: PLC0415
+
+        if parent.metadata is None:
+            raise ValueError(
+                "can't compute filesystem path for document withtout metadata"
+            )
+        if parent.is_root_document:
+            path = parent.path
+            if path is None:
+                msg = (
+                    "can't create a Datashare file from ES document without path, "
+                    "path is needed to locate the file, retrieve it from ES"
+                )
+                raise ValueError(msg)
+            location = FileLocation.FILESYSTEM
+        else:
+            path = artifacts_dir(doc_id=parent.id, project=parent.project) / "raw"
+            location = FileLocation.ARTIFACTS
+        # The filesystem dod is alway relative to the base location, let's make sure
+        # 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:])
+        n_pages = 1
+        if parent.metadata:
+            n_pages = parent.metadata.get("tika_metadata_xmptpg_npages", n_pages)
+        path = WorkerFilePath(path=path, location=location)
+        return cls(
+            doc_id=parent.id,
+            parent=parent,
+            value=path,
+            n_pages=n_pages,
+            project=parent.project,
+        )
+
+
+class WorkerFile(
+    WithParent["ProcessedFile"], ProcessingResult[WorkerFilePath], ErrorSourceWithPages
+):
+    """An intermediate results produce by the worker either from a Datashare doc on the
+    filesystem or from another intermediate result
+    """
+
+    parent: "ProcessedFile"
+
+    def locate(self, paths: WorkerPaths) -> Path:
+        return self.value.locate(paths)
+
+    @classmethod
+    def from_parent(
+        cls, parent: "ProcessedFile", path: Path, paths: WorkerPaths
+    ) -> Self:
+        path = WorkerFilePath.relative_to_workdir(path, paths)
+        return cls(
+            doc_id=parent.doc_id,
+            project=parent.project,
+            parent=parent,
+            value=path,
+            n_pages=parent.n_pages,
+        )
+
+
+def _file_discriminator(v: Any) -> str:
+    parent = v.get("parent") if isinstance(v, dict) else getattr(v, "parent", None)
+    if parent is None:
+        raise ValueError(f"{v} has no parent")
+    if isinstance(parent, dict):
+        is_document = parent.get("type") == "Document"
+    else:
+        is_document = isinstance(parent, Document)
+    return "datashare" if is_document else "worker"
+
+
+ProcessedFile = Annotated[
+    Annotated[DatashareFile, Tag("datashare")] | Annotated[WorkerFile, Tag("worker")],
+    Discriminator(_file_discriminator),
+]
+WorkerFile.model_rebuild()
+PROCESSED_FILE_TA = TypeAdapter(ProcessedFile, config=BaseModel.model_config)

{datashare_python-0.11.3 → datashare_python-0.12.0}/datashare_python/utils.py

@@ -10,6 +10,7 @@ import sys
 import threading
 import time
 from collections.abc import (
+    AsyncGenerator,
     AsyncIterable,
     Awaitable,
     Callable,
@@ -25,13 +26,13 @@ from functools import cache, wraps
 from hashlib import sha256
 from io import BytesIO
 from pathlib import Path
-from typing import Any, Self, TypeVar
+from typing import Any, Protocol, Self, TypeVar
 from uuid import uuid4
 
 import temporalio
 from aiofile import async_open
 from lru import LRU
-from pydantic import ValidationError
+from pydantic import TypeAdapter, ValidationError
 from pydantic.alias_generators import to_snake
 from temporalio import activity, workflow
 from temporalio.api.common.v1 import Payload
@@ -56,9 +57,13 @@ from .constants import MANIFEST_JSON, METADATA_JSON
 from .mimetypes_ import types_map
 from .objects import (
     BaseModel,
+    DatashareFile,
     DocArtifact,
-    DocumentLocation,
+    Document,
+    ErrorSource,
     ProcessedFile,
+    ProcessingError,
+    WorkerFile,
     WorkerPaths,
 )
 from .types_ import RawAsyncProgressHandler
@@ -388,8 +393,7 @@ def to_scaled_async_progress(
 def safe_dir(doc_id: str) -> Path:
     if len(doc_id) < 4:
         raise ValueError(f"expected doc_id to be at least 4, found {doc_id}")
-    parts = (p for p in (doc_id[:2], doc_id[2:4]) if p)
-    return Path(*parts)
+    return Path(doc_id[:2], doc_id[2:4])
 
 
 def artifacts_dir(doc_id: str, *, project: str) -> Path:
@@ -607,31 +611,18 @@ def activity_workdir(
 
 
 def symlink_embedded_document_to_workdir(
-    doc: ProcessedFile, paths: WorkerPaths
+    doc: Document, paths: WorkerPaths
 ) -> ProcessedFile:
-    match doc.location:
-        case DocumentLocation.ARTIFACTS:
-            symlinks_dir = paths.workdir / doc.project / "symlinks"
-            symlinks_dir.mkdir(parents=True, exist_ok=True)
-            symlink_path = Path(*doc.path.parts[1:-1], doc.resource_name)
-            symlink_path = symlinks_dir / symlink_path
-            # Replace the "raw" with the resource name
-            symlink_path.parent.mkdir(parents=True, exist_ok=True)
-            artifact_path = paths.artifacts / doc.path
-            with contextlib.suppress(FileExistsError):
-                os.symlink(artifact_path, symlink_path)
-            return ProcessedFile(
-                path=symlink_path.relative_to(paths.workdir),
-                id=doc.id,
-                location=DocumentLocation.WORKDIR,
-                project=doc.project,
-                resource_name=doc.resource_name,
-                n_pages=doc.n_pages,
-            )
-        case DocumentLocation.FILESYSTEM:
-            return doc
-        case _:
-            raise ValueError(f"unsupported location {doc.location}")
+    ds_file = DatashareFile.from_parent(doc)
+    if doc.is_root_document:
+        return ds_file
+    symlinks_dir = paths.workdir / doc.project / "symlinks"
+    symlink_path = symlinks_dir / safe_dir(doc.id) / doc.id / doc.resource_name
+    symlink_path.parent.mkdir(parents=True, exist_ok=True)
+    raw_path = ds_file.locate(paths)
+    with contextlib.suppress(FileExistsError):
+        os.symlink(raw_path, symlink_path)
+    return WorkerFile.from_parent(ds_file, symlink_path, paths)
 
 
 def artifact_path(
@@ -644,7 +635,7 @@ def artifact_path(
     return root / artifacts_dir(doc_id, project=project) / artifact_type.filename
 
 
-def _read_jsonl(path: Path) -> Iterable[dict]:
+def read_jsonl(path: Path) -> Iterable[dict]:
     with path.open() as f:
         for line in f:
             line = line.strip()  # noqa: PLW2901
@@ -655,18 +646,82 @@ def _read_jsonl(path: Path) -> Iterable[dict]:
 M = TypeVar("M", bound=BaseModel)
 
 
-def read_jsonl_as[M](path: Path, cls: type[M]) -> Iterable[M]:
-    return (cls.model_validate(d) for d in _read_jsonl(path))
+def read_jsonl_as[M](
+    path: Path, processed_file_cls: type[M] | TypeAdapter[M]
+) -> Iterable[M]:
+    if isinstance(processed_file_cls, TypeAdapter):
+        deser_fn = processed_file_cls.validate_python
+    else:
+        deser_fn = processed_file_cls.model_validate
+    return (deser_fn(d) for d in read_jsonl(path))
 
 
 async def async_read_jsonl_as[M](
-    path: Path, processed_file_cls: type[M]
+    path: Path, processed_file_cls: type[M] | TypeAdapter[M]
 ) -> AsyncIterable[M]:
+    if isinstance(processed_file_cls, TypeAdapter):
+        deser_fn = processed_file_cls.validate_json
+    else:
+        deser_fn = processed_file_cls.model_validate_json
     async with async_open(path, "r") as f:
         async for line in f:
             line = line.strip()  # noqa: PLW2901
             if line:
-                yield processed_file_cls.model_validate_json(line)
+                yield deser_fn(line)
+
+
+async def write_batches(
+    batches: AsyncIterable[list[ProcessedFile]] | Iterable[list[ProcessedFile]],
+    root: Path,
+    batch_offset: int = 0,
+    prefix: str = "batch_",
+) -> AsyncIterable[Path]:
+    if hasattr(batches, "__aiter__"):
+        async for b in _async_write_batches(batches, root, batch_offset, prefix):
+            yield b
+        return
+    async for b in _write_batches(batches, root, batch_offset, prefix):
+        yield b
+
+
+async def _async_write_batches(
+    batches: AsyncIterable[list[ProcessedFile]],
+    root: Path,
+    batch_offset: int,
+    prefix: str,
+) -> AsyncIterable[Path]:
+    batch_id = batch_offset
+    async for batch in batches:
+        batch_path = root / f"{batch_id // 1000}" / f"{prefix}{batch_id}.jsonl"
+        batch_path.parent.mkdir(parents=True, exist_ok=True)
+        async with async_open(batch_path, "w") as f:
+            for fs_doc in batch:
+                serialized = fs_doc.model_dump_json(
+                    exclude_none=True, polymorphic_serialization=True
+                )
+                await f.write(f"{serialized}\n")
+        yield batch_path
+        batch_id += 1
+
+
+async def _write_batches(
+    batches: Iterable[list[ProcessedFile]],
+    root: Path,
+    batch_offset: int,
+    prefix: str,
+) -> AsyncIterable[Path]:
+    batch_id = batch_offset
+    for batch in batches:
+        batch_path = root / f"{batch_id // 1000}" / f"{prefix}{batch_id}.jsonl"
+        batch_path.parent.mkdir(parents=True, exist_ok=True)
+        async with async_open(batch_path, "w") as f:
+            for fs_doc in batch:
+                serialized = fs_doc.model_dump_json(
+                    exclude_none=True, polymorphic_serialization=True
+                )
+                await f.write(f"{serialized}\n")
+        yield batch_path
+        batch_id += 1
 
 
 @cache
@@ -689,6 +744,8 @@ def ext_to_mime_types(ext: str) -> set[str]:
             return ext_to_mime_types(".xml")
         case ".dclg.xml":
             return ext_to_mime_types(".xml")
+        case ".msg":
+            return ext_to_mime_types(".eml")
     try:
         return types_map()[ext]
     except KeyError as e:
@@ -766,59 +823,108 @@ class SharedResources:
         eviction_callback: Callable[[Any, Any], None] | None = None,
     ) -> None:
         self._eviction_callback = eviction_callback
-        self._cache = LRU(cache_size, self._eviction_callback)
+        self._cache = LRU(cache_size, self._on_evict)
         self._sentinel = object()
+        # Leased resources are tracked by id, evicting a leased resource defers the
+        # eviction callback until its last lease is released
+        self._lock = threading.Lock()
+        self._n_leases: dict[int, int] = dict()
+        self._pending_evictions: dict[int, tuple[Any, Any]] = dict()
 
     def __enter__(self) -> Self:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:  # noqa: ANN001
-        if self._eviction_callback is not None:
-            for k, v in self._cache.items():
-                self._eviction_callback(k, v)
+        for k, v in self._cache.items():
+            self._on_evict(k, v)
 
+    @contextlib.contextmanager
     def get_or_cache_resource(
         self, key: str, default_factory: Callable[[], Any]
-    ) -> Any:
-        value = self._cache.get(key, self._sentinel)
-        if value is not self._sentinel:
-            return value
-        # Get the value first to be sure to return the right one in case of concurrent
-        # access.
-        #
-        # Additionally, the factory can be long to complete (that's one of the reason
-        # we want to cache its results). Because we don't lock the factory call,
-        # together with the cache update, in case of concurrent access,
-        # quicker factories will complete first and will stay longer in the cache.
-        # This is fine, the alternative is to lock factory + cache update to avoid
-        # concurrent access but this will mean waiting for the first factory call to
-        # complete, this can potentially imply longer waits than no getting the value
-        # from the cache
-        value = default_factory()
-        self._cache[key] = value
-        return value
+    ) -> Generator[Any, None, None]:
+        """Get or cache a resource, leasing it for the duration of the context.
+
+        A resource evicted while leased stays usable, its eviction callback only runs
+        once all its leases have been released.
+
+        :param key: resource cache key
+        :param default_factory: resource factory, called on cache miss
+        :return: the cached resource
+        """
+        value = self._get_and_acquire(key)
+        if value is self._sentinel:
+            value = default_factory()
+            self._cache_and_acquire(key, value)
+        try:
+            yield value
+        finally:
+            self._release(value)
 
+    @contextlib.asynccontextmanager
     async def async_get_or_cache_resource(
         self, key: str, default_factory: Callable[[], Awaitable[Any]]
-    ) -> Any:
-        value = self._cache.get(key, self._sentinel)
-        if value is not self._sentinel:
-            return value
-        # Get the value first to be sure to return the right one in case of concurrent
-        # access.
-        #
-        # Additionally, the factory can be long to complete (that's one of the reason
-        # we want to cache its results). Because we don't lock the factory call,
-        # together with the cache update, in case of concurrent access,
-        # quicker factories will complete first and will stay longer in the cache.
-        # This is fine, the alternative is to lock factory + cache update to avoid
-        # concurrent access but this will mean waiting for the first factory call to
-        # complete, this can potentially imply longer waits than no getting the value
-        # from the cache
-        value = await default_factory()
-        self._cache[key] = value
+    ) -> AsyncGenerator[Any, None]:
+        """Async version of get_or_cache_resource.
+
+        :param key: resource cache key
+        :param default_factory: async resource factory, awaited on cache miss
+        :return: the cached resource
+        """
+        value = self._get_and_acquire(key)
+        if value is self._sentinel:
+            value = await default_factory()
+            self._cache_and_acquire(key, value)
+        try:
+            yield value
+        finally:
+            self._release(value)
+
+    def _get_and_acquire(self, key: str) -> Any:
+        with self._lock:
+            value = self._cache.get(key, self._sentinel)
+            if value is not self._sentinel:
+                self._acquire(value)
         return value
 
+    def _cache_and_acquire(self, key: str, value: Any) -> None:
+        # The factory can be long to complete (that's one of the reason we want to
+        # cache its results). Because we don't lock the factory call, together with
+        # the cache update, in case of concurrent access, quicker factories will
+        # complete first and will stay longer in the cache. This is fine, the
+        # alternative is to lock factory + cache update to avoid concurrent access but
+        # this will mean waiting for the first factory call to complete, this can
+        # potentially imply longer waits than no getting the value from the cache
+        with self._lock:
+            self._acquire(value)
+        # Can evict other resources, must be called without holding the lock
+        self._cache[key] = value
+
+    def _acquire(self, value: Any) -> None:
+        value_id = id(value)
+        self._n_leases[value_id] = self._n_leases.get(value_id, 0) + 1
+
+    def _release(self, value: Any) -> None:
+        value_id = id(value)
+        with self._lock:
+            self._n_leases[value_id] -= 1
+            if self._n_leases[value_id]:
+                return
+            del self._n_leases[value_id]
+            pending = self._pending_evictions.pop(value_id, None)
+        if pending is not None:
+            self._evict(*pending)
+
+    def _on_evict(self, key: Any, value: Any) -> None:
+        with self._lock:
+            if id(value) in self._n_leases:
+                self._pending_evictions[id(value)] = (key, value)
+                return
+        self._evict(key, value)
+
+    def _evict(self, key: Any, value: Any) -> None:
+        if self._eviction_callback is not None:
+            self._eviction_callback(key, value)
+
 
 def close_cm_callback(key: str, value: Any) -> None:  # noqa: ARG001
     if hasattr(value, "__exit__"):
@@ -857,3 +963,43 @@ def config_cache_key(config: BaseModel) -> str:
         .replace("_", "-")
     )
     return f"{config_label}-{str(hash(config))}"
+
+
+class _PreprocessingFunction[S: ErrorSource, R](Protocol):
+    def __call__(self, doc: S, *args, **kwargs) -> R: ...
+
+
+def reports_errors[S: ErrorSource, R](
+    errors: tuple[type[Exception]], exc_cls: type[ProcessingError[S]]
+) -> Callable[
+    [_PreprocessingFunction[S, R]], _PreprocessingFunction[S, R | ProcessingError[S]]
+]:
+
+    def parent_wrapper(
+        f: _PreprocessingFunction[S, R],
+    ) -> _PreprocessingFunction[S, R | ProcessingError[S]]:
+        if inspect.iscoroutinefunction(f):
+
+            @wraps(f)  # noqa: F821
+            async def wrapper(doc: S, *args, **kwargs) -> R | ProcessingError[S]:
+                try:
+                    return await f(doc, *args, **kwargs)
+                except errors as e:
+                    logger.exception("error while processing doc %s", doc)
+                    report = exc_cls.from_exception(doc, e)
+                    return report
+        else:
+
+            @wraps(f)
+            def wrapper(doc: ProcessedFile, *args, **kwargs) -> R | ProcessingError[S]:
+
+                try:
+                    return f(doc, *args, **kwargs)
+                except errors as e:
+                    logger.exception("error while processing doc %s", doc)
+                    report = exc_cls.from_exception(doc, e)
+                    return report
+
+        return wrapper
+
+    return parent_wrapper

{datashare_python-0.11.3 → datashare_python-0.12.0}/pyproject.toml

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