datashare-python 0.8.19__tar.gz → 0.8.20__tar.gz

{datashare_python-0.8.19 → datashare_python-0.8.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.8.19
+Version: 0.8.20
 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/
@@ -12,8 +12,10 @@ Requires-Dist: aiohttp~=3.11
 Requires-Dist: alive-progress~=3.2
 Requires-Dist: hatchling~=1.27
 Requires-Dist: icij-common[elasticsearch]~=0.8.2
+Requires-Dist: langcodes~=3.5
 Requires-Dist: nest-asyncio~=1.6
 Requires-Dist: orjson~=3.11
+Requires-Dist: pydantic-extra-types[pycountry]>=2.11.1
 Requires-Dist: python-json-logger~=4.0
 Requires-Dist: pyyaml~=6.0
 Requires-Dist: temporalio~=1.23

{datashare_python-0.8.19 → datashare_python-0.8.20}/datashare_python/conftest.py

@@ -1,4 +1,5 @@
 import asyncio
+import shutil
 from asyncio import AbstractEventLoop
 from collections.abc import AsyncGenerator, Generator, Iterator, Sequence
 from pathlib import Path
@@ -53,7 +54,6 @@ _INDEX_BODY = {
             "join": {"type": "join", "relations": {"Document": "NamedEntity"}},
             "contentType": {"type": "keyword"},
             "content": {"type": "text"},
-            "contentTranslated": {"type": "text"},
         }
     }
 }
@@ -197,10 +197,7 @@ def index_docs_ops(
     docs: list[Document], index_name: str
 ) -> Generator[dict, None, None]:
     for doc in docs:
-        op = {
-            "_op_type": "index",
-            "_index": index_name,
-        }
+        op = {"_op_type": "index", "_index": index_name}
         doc = doc.model_dump(by_alias=True)  # noqa: PLW2901
         op.update(doc)
         if "path" in op:
@@ -305,3 +302,10 @@ async def all_done(task_client: DatashareTaskClient, not_done: list[str]) -> boo
 @pytest.fixture  # noqa: F405
 def typer_asyncio_patch() -> None:
     nest_asyncio.apply()
+
+
+def clear_dirs(config: WorkerConfig) -> None:
+    shutil.rmtree(str(config.artifacts_root))
+    config.artifacts_root.mkdir(parents=True, exist_ok=True)
+    shutil.rmtree(str(config.workdir))
+    config.workdir.mkdir(parents=True, exist_ok=True)

{datashare_python-0.8.19 → datashare_python-0.8.20}/datashare_python/objects.py

@@ -6,8 +6,12 @@ from datetime import UTC, datetime
 from enum import StrEnum, unique
 from io import BytesIO
 from pathlib import Path
-from typing import Annotated, Any, Literal, Self, TypeVar, cast
+from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, cast
 
+import langcodes
+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 TIKA_METADATA_RESOURCENAME
@@ -31,7 +35,14 @@ from icij_common.pydantic_utils import (
     merge_configs,
     no_enum_values_config,
 )
-from pydantic import AfterValidator, Field
+from pydantic import (
+    AfterValidator,
+    BeforeValidator,
+    Field,
+    GetCoreSchemaHandler,
+    TypeAdapter,
+    model_validator,
+)
 from pydantic import BaseModel as _BaseModel
 from pydantic.main import IncEx
 
@@ -50,110 +61,45 @@ class DatashareModel(BaseModel):
     model_config = merge_configs(BaseModel.model_config, lowercamel_case_config())
 
 
-@unique
-class TaskState(StrEnum):
-    CREATED = "CREATED"
-    QUEUED = "QUEUED"
-    RUNNING = "RUNNING"
-    ERROR = "ERROR"
-    DONE = "DONE"
-    CANCELLED = "CANCELLED"
-
-
-READY_STATES = frozenset({TaskState.DONE, TaskState.ERROR, TaskState.CANCELLED})
+class DatashareLanguage(str):
+    _language_type_adapter: ClassVar[TypeAdapter] = TypeAdapter(LanguageName)
 
+    @classmethod
+    def _validate(cls, __input_value: str, _: core_schema.ValidationInfo) -> Self:
+        if __input_value != __input_value.upper():
+            raise PydanticCustomError(
+                "datashare_language", "Invalid Datashare language, expected uppercase"
+            )
+        try:
+            # Use pydantic provided validation
+            cls._language_type_adapter.validate_python(__input_value.title())
+        except ValidationError as e:
+            raise PydanticCustomError(
+                "datashare_language", "Unknown Datashare language"
+            ) from e
+        return cls(__input_value)
 
-class StacktraceItem(DatashareModel):
-    name: str
-    file: str
-    lineno: int
-
-
-class Message(DatashareModel):
-    type: str = Field(frozen=True, alias="@type")
-
-    def model_dump(
-        self,
-        *,
-        mode: Literal["json", "python"] | str = "python",
-        include: IncEx | None = None,
-        exclude: IncEx | None = None,
-        context: Any | None = None,
-        exclude_unset: bool = False,
-        exclude_defaults: bool = False,
-        exclude_none: bool = False,
-        round_trip: bool = False,
-        warnings: bool | Literal["none", "warn", "error"] = True,
-        fallback: Callable[[Any], Any] | None = None,
-        serialize_as_any: bool = False,
-    ) -> dict[str, Any]:
-        return super().model_dump(
-            by_alias=True,
-            mode=mode,
-            include=include,
-            exclude=exclude,
-            context=context,
-            exclude_unset=exclude_unset,
-            exclude_defaults=exclude_defaults,
-            exclude_none=exclude_none,
-            round_trip=round_trip,
-            warnings=warnings,
-            fallback=fallback,
-            serialize_as_any=serialize_as_any,
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source: type[Any], handler: GetCoreSchemaHandler
+    ) -> core_schema.AfterValidatorFunctionSchema:
+        return core_schema.with_info_after_validator_function(
+            cls._validate,
+            core_schema.str_schema(),
+            serialization=core_schema.to_string_ser_schema(),
         )
 
+    @property
+    def as_language_name(self) -> LanguageName:
+        return LanguageName(self.title())
 
-class TaskResult(Message):
-    type: str = Field(frozen=True, alias="@type", default="TaskResult")
-    value: object
-
-
-class TaskError(Message):
-    type: str = Field(frozen=True, alias="@type", default="TaskError")
-    name: str
-    message: str
-    cause: str | None = None
-    stacktrace: list[StacktraceItem] = Field(default_factory=list)
-
-
-def _datetime_now() -> datetime:
-    return datetime.now(UTC)
-
-
-class User(Message):
-    type: str = Field(
-        frozen=True, alias="@type", default="org.icij.datashare.user.User"
-    )
-    id: str
-    name: str | None = None
-    email: str | None = None
-    provider: str | None = None
-    details: dict = dict()
-
-
-class Task(Message):
-    type: str = Field(frozen=True, alias="@type", default="Task")
-    id: str
-    name: str
-    args: dict[str, object] | None = None
-    state: TaskState = TaskState.CREATED
-    result: TaskResult | None = None
-    error: TaskError | None = None
-    progress: float | None = None
-    created_at: datetime = Field(default_factory=_datetime_now)
-    completed_at: datetime | None = None
-    retries_left: int | None = None
-    max_retries: int | None = None
-
-
-@dataclass(frozen=True)
-class TaskGroup:
-    name: str
+    @property
+    def alpha2(self) -> str | None:
+        return self.as_language_name.alpha2
 
     @property
-    @classmethod
-    def python(cls) -> Self:
-        return cls(name="PYTHON")
+    def alpha3(self) -> str:
+        return self.as_language_name.alpha3
 
 
 @unique
@@ -195,21 +141,64 @@ class FilesystemDocument(DatashareModel):
                 raise ValueError(f"invalid location: {self.path}")
 
 
+class IETFLanguage(str):
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source: Any, handler: GetCoreSchemaHandler
+    ) -> PlainValidatorFunctionSchema:
+        return core_schema.no_info_plain_validator_function(cls.validate)
+
+    @classmethod
+    def validate(cls, v: Any) -> Self:
+        tag = langcodes.get(str(v))
+        if not tag.is_valid():
+            raise ValueError(f"Invalid IETF language: {v}")
+        return cls(v)
+
+
+Language = DatashareLanguage | IETFLanguage
+
+
+def _from_sentences(value: Any) -> Any:
+    if isinstance(value, list):
+        return " ".join(value)
+    return value
+
+
+class Translation(BaseModel):  # No camelcase here we don't know why
+    source_language: DatashareLanguage
+    target_language: Language
+    translator: str
+    content: Annotated[str, BeforeValidator(_from_sentences)]
+
+
 class Document(DatashareModel):
     id: str
-    language: str
+    language: DatashareLanguage
     index: str | None = None
     root_document: str | None = None
     content: str | None = None
+    content_text_length: int | None = None
     content_type: str | None = None
     path: Path | None = None
     tags: list[str] = Field(default_factory=list)
-    content_translated: dict[str, str] = Field(
-        default_factory=dict, alias="content_translated"
+    content_translated: list[Translation] | None = Field(
+        default=None,
+        # es translator is using snake_case, we must do the same
+        alias="content_translated",
     )
     metadata: dict[str, Any] | None = None
     type: str = Field(default="Document", frozen=True)
 
+    @model_validator(mode="before")
+    @classmethod
+    def _initialize_content_length_from_content(cls, data: Any) -> Any:
+        if isinstance(data, dict):
+            content_length = data.get("content_text_length")
+            if content_length is None and (content := data.get(DOC_CONTENT)):
+                data["content_text_length"] = len(content)
+        return data
+
     @classmethod
     def from_es(cls, es_doc: dict) -> Self:
         sources = es_doc[SOURCE]
@@ -217,8 +206,9 @@ class Document(DatashareModel):
             id=es_doc[ID_],
             index=es_doc.get(INDEX_),
             content=sources.get(DOC_CONTENT),
-            content_translated=sources.get(DOC_CONTENT_TRANSLATED, dict()),
-            language=sources[DOC_LANGUAGE],
+            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),
             tags=sources.get("tags", []),
             path=sources.get(DOC_PATH),
@@ -258,10 +248,122 @@ class Document(DatashareModel):
         )
 
 
+def _is_absolute_path(v: bytes | BytesIO | Path) -> Any:
+    if isinstance(v, Path) and not v.is_absolute():
+        raise ValueError("artifact path must be absolute")
+    return v
+
+
 @dataclass(frozen=True)
 class DocArtifact:
     project: str
     doc_id: str
-    artifact: bytes | BytesIO
+    artifact: Annotated[bytes | BytesIO | Path, AfterValidator(_is_absolute_path)]
     filename: str
     metadata_key: str
+
+
+@unique
+class TaskState(StrEnum):
+    CREATED = "CREATED"
+    QUEUED = "QUEUED"
+    RUNNING = "RUNNING"
+    ERROR = "ERROR"
+    DONE = "DONE"
+    CANCELLED = "CANCELLED"
+
+
+READY_STATES = frozenset({TaskState.DONE, TaskState.ERROR, TaskState.CANCELLED})
+
+
+class StacktraceItem(DatashareModel):
+    name: str
+    file: str
+    lineno: int
+
+
+class Message(DatashareModel):
+    type: str = Field(frozen=True, alias="@type")
+
+    def model_dump(
+        self,
+        *,
+        mode: Literal["json", "python"] | str = "python",
+        include: IncEx | None = None,
+        exclude: IncEx | None = None,
+        context: Any | None = None,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        round_trip: bool = False,
+        warnings: bool | Literal["none", "warn", "error"] = True,
+        fallback: Callable[[Any], Any] | None = None,
+        serialize_as_any: bool = False,
+    ) -> dict[str, Any]:
+        return super().model_dump(
+            by_alias=True,
+            mode=mode,
+            include=include,
+            exclude=exclude,
+            context=context,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+            fallback=fallback,
+            serialize_as_any=serialize_as_any,
+        )
+
+
+class TaskResult(Message):
+    type: str = Field(frozen=True, alias="@type", default="TaskResult")
+    value: object
+
+
+class TaskError(Message):
+    type: str = Field(frozen=True, alias="@type", default="TaskError")
+    name: str
+    message: str
+    cause: str | None = None
+    stacktrace: list[StacktraceItem] = Field(default_factory=list)
+
+
+def _datetime_now() -> datetime:
+    return datetime.now(UTC)
+
+
+class User(Message):
+    type: str = Field(
+        frozen=True, alias="@type", default="org.icij.datashare.user.User"
+    )
+    id: str
+    name: str | None = None
+    email: str | None = None
+    provider: str | None = None
+    details: dict = dict()
+
+
+class Task(Message):
+    type: str = Field(frozen=True, alias="@type", default="Task")
+    id: str
+    name: str
+    args: dict[str, object] | None = None
+    state: TaskState = TaskState.CREATED
+    result: TaskResult | None = None
+    error: TaskError | None = None
+    progress: float | None = None
+    created_at: datetime = Field(default_factory=_datetime_now)
+    completed_at: datetime | None = None
+    retries_left: int | None = None
+    max_retries: int | None = None
+
+
+@dataclass(frozen=True)
+class TaskGroup:
+    name: str
+
+    @property
+    @classmethod
+    def python(cls) -> Self:
+        return cls(name="PYTHON")

{datashare_python-0.8.19 → datashare_python-0.8.20}/datashare_python/template.py

@@ -41,7 +41,8 @@ def build_template_tarball() -> None:
             is_hidden = path.name.startswith(".") or any(
                 "." in p for p in path.parts[:-1]
             )
-            if is_hidden or not path.is_file() or path.suffix not in ALLOWED_EXTS:
+            skip = is_hidden or not path.is_file() or path.suffix not in ALLOWED_EXTS
+            if path.name != "Dockerfile" and skip:
                 continue
             tar.add(path, arcname=path.relative_to(template_dir))
 
@@ -75,8 +76,11 @@ def _update_pyproject_toml(
     pyproject_toml["tool"]["uv"].pop("index", None)
 
     project = pyproject_toml["project"]
+    project["name"] = package_name.replace("_", "-").lower()
+    project["version"] = "0.1.0"
     project["authors"] = []
     project.pop("urls", None)
+    project.pop("description", None)
     project["dependencies"] = sorted(
         d
         for d in project["dependencies"]
@@ -100,16 +104,29 @@ def _update_pyproject_toml(
         "worker_template", package_name
     )
     entry_points["datashare.activities"]["activities"] = activities_entry_point
+
+    deps = entry_points["datashare.dependencies"]["dependencies"]
+    deps = deps.replace("worker_template", package_name)
+    entry_points["datashare.dependencies"]["dependencies"] = deps
+
+    cfg_cls_entry_point = entry_points["datashare.worker_config_cls"][
+        "worker_config_cls"
+    ]
+    cfg_cls_entry_point = cfg_cls_entry_point.replace("worker_template", package_name)
+    entry_points["datashare.worker_config_cls"]["worker_config_cls"] = (
+        cfg_cls_entry_point
+    )
+
     hatch_sdist = pyproject_toml["tool"]["hatch"]["build"]["targets"]["wheel"]
     hatch_sdist["packages"] = [
         i if i != "worker_template" else package_name for i in hatch_sdist["packages"]
     ]
 
-    hatch_sdist = pyproject_toml["tool"]["hatch"]["build"]["targets"]["sdist"]
-    if "only-include" in hatch_sdist:
-        hatch_sdist["only-include"] = [
-            i if i != "worker_template" else package_name
-            for i in hatch_sdist["only-include"]
-        ]
+    hatch_sdist = pyproject_toml["tool"]["hatch"]["build"]["targets"]["wheel"]
+    if "packages" in hatch_sdist:
+        hatch_sdist["packages"] = [package_name]
+
+    build_system = pyproject_toml["build-system"]
+    build_system["package"] = package_name
 
     return pyproject_toml

{datashare_python-0.8.19 → datashare_python-0.8.20}/datashare_python/utils.py

@@ -3,8 +3,10 @@ import contextlib
 import contextvars
 import inspect
 import json
+import os
+import shutil
 import threading
-from collections.abc import Awaitable, Callable, Coroutine
+from collections.abc import Awaitable, Callable, Coroutine, Iterable
 from copy import deepcopy
 from dataclasses import dataclass
 from datetime import timedelta
@@ -24,7 +26,7 @@ from temporalio.common import RetryPolicy, SearchAttributeKey
 from temporalio.exceptions import ApplicationError
 
 from .constants import METADATA_JSON
-from .objects import DocArtifact
+from .objects import DocArtifact, DocumentLocation, FilesystemDocument
 from .types_ import ProgressRateHandler, RawProgressHandler
 
 DependencyLabel = str | None
@@ -83,7 +85,6 @@ class WorkflowWithProgress:
     async def update_progress(self, signal: ProgressSignal) -> None:
         async with self._update_lock:
             # TODO: remove this log
-            workflow.logger.debug("recording progress signal %s", signal)
             key = (signal.run_id, signal.activity_id)
             self._progress[key] = signal.to_progress()
             progress = sum(p.current for p in self._progress.values())
@@ -480,11 +481,17 @@ def write_artifact(root: Path, artifact: DocArtifact) -> Path:
     # TODO: if transcriptions are too large we could also serialize them
     #  as jsonl
     artifact_path: Path = artif_dir / artifact.filename
-    if isinstance(artifact.artifact, bytes):
-        artifact_path.write_bytes(artifact.artifact)
-    elif isinstance(artifact_path, BytesIO):
-        with artifact_path.open("wb") as f:
-            f.write(artifact.artifact.read())
+    match artifact.artifact:
+        case bytes():
+            artifact_path.write_bytes(artifact.artifact)
+        case BytesIO():
+            with artifact_path.open("wb") as f:
+                f.write(artifact.artifact.read())
+        case Path():
+            shutil.move(artifact.artifact, artifact_path)
+        case _:
+            msg = f"unsupported artifact type: {artifact.artifact.__class__.__name__}"
+            raise ValueError(msg)
     meta_path = root / _metadata_path(artifact.doc_id, project=artifact.project)
     meta = _read_artifact_metadata(root, artifact) if meta_path.exists() else dict()
     meta[artifact.metadata_key] = artifact.filename
@@ -557,3 +564,40 @@ def activity_workdir(
         wf_context=wf_context, act_context=act_context, run_context=run_context
     )
     return workdir.joinpath(project, ctx_path)
+
+
+def symlink_embedded_document_to_workdir(
+    doc: FilesystemDocument, artifacts_root: Path, *, workdir: Path
+) -> FilesystemDocument:
+    match doc.location:
+        case DocumentLocation.ARTIFACTS:
+            symlinks_dir = workdir / doc.index / "symlinks"
+            symlinks_dir.mkdir(parents=True, exist_ok=True)
+            symlink_path = Path(*doc.path.parts[:-1], doc.id)
+            # Replace the "raw" with the doc id
+            doc_ext = Path(doc.resource_name).suffix
+            symlink_path = symlink_path.relative_to(Path(doc.index))
+            symlink_path = symlinks_dir / f"{symlink_path}{doc_ext}"
+            symlink_path.parent.mkdir(parents=True, exist_ok=True)
+            artifact_path = artifacts_root / doc.path
+            with contextlib.suppress(FileExistsError):
+                os.symlink(artifact_path, symlink_path)
+            return FilesystemDocument(
+                path=symlink_path.relative_to(workdir),
+                id=doc.id,
+                location=DocumentLocation.WORKDIR,
+                index=doc.index,
+                resource_name=doc.resource_name,
+            )
+        case DocumentLocation.ORIGINAL:
+            return doc
+        case _:
+            raise ValueError(f"unsupported location {doc.location}")
+
+
+def read_jsonl(path: Path) -> Iterable[dict]:
+    with path.open() as f:
+        for line in f:
+            line = line.strip()  # noqa: PLW2901
+            if line:
+                yield json.loads(line)

{datashare_python-0.8.19 → datashare_python-0.8.20}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datashare-python"
-version = "0.8.19"
+version = "0.8.20"
 description = "Manage Python tasks and local resources in Datashare"
 authors = [
   { name = "Clément Doumouro", email = "cdoumouro@icij.org" },
@@ -12,15 +12,17 @@ requires-python = ">=3.11,<4"
 dependencies = [
   "alive-progress~=3.2",
   "aiohttp~=3.11",
+  "hatchling~=1.27",
   "icij-common[elasticsearch]~=0.8.2",
-  "python-json-logger~=4.0",
+  "langcodes~=3.5",
   "nest-asyncio~=1.6",
+  "orjson~=3.11",
+  "python-json-logger~=4.0",
+  "pyyaml~=6.0",
+  "pydantic-extra-types[pycountry]>=2.11.1",
   "temporalio~=1.23",
   "typer>=0.15.4,<0.25.1",
   "tomlkit~=0.14.0",
-  "hatchling~=1.27",
-  "pyyaml~=6.0",
-  "orjson~=3.11",
 ]
 
 [project.urls]