datashare-python 0.2.23__tar.gz → 0.2.25__tar.gz

{datashare_python-0.2.23 → datashare_python-0.2.25}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.2.23
+Version: 0.2.25
 Summary: Manage Pythoœn tasks and local resources in Datashare
 Project-URL: Homepage, https://icij.github.io/datashare-python/
 Project-URL: Documentation, https://icij.github.io/datashare-python/
@@ -9,13 +9,13 @@ Project-URL: Issues, https://github.com/ICIJ/datashare-python/issues
 Author-email: Clément Doumouro <cdoumouro@icij.org>, Clément Doumouro <clement.doumouro@gmail.com>, Lion Summerbell <lsummerbell@icij.org>
 Requires-Python: <4,>=3.11
 Requires-Dist: aiohttp~=3.11.9
-Requires-Dist: aiostream~=0.6.4
 Requires-Dist: alive-progress~=3.2.0
-Requires-Dist: datashare-worker-template~=0.1
+Requires-Dist: datashare-worker-template[ml]~=0.1
 Requires-Dist: hatchling~=1.27.0
-Requires-Dist: icij-common[elasticsearch]~=0.7.3
+Requires-Dist: icij-common[elasticsearch]~=0.8.2
 Requires-Dist: nest-asyncio~=1.6.0
 Requires-Dist: python-json-logger~=4.0.0
+Requires-Dist: pyyaml~=6.0
 Requires-Dist: temporalio~=1.23.0
 Requires-Dist: tomlkit~=0.14.0
 Requires-Dist: typer~=0.15.4

datashare_python-0.2.25/datashare_python/cli/worker.py

@@ -0,0 +1,137 @@
+import asyncio
+import logging
+from pathlib import Path
+from typing import Annotated
+
+import typer
+import yaml
+from icij_common.pydantic_utils import safe_copy
+
+from datashare_python.config import TemporalClientConfig, WorkerConfig
+from datashare_python.discovery import discover, discover_activities, discover_workflows
+from datashare_python.worker import bootstrap_worker, create_worker_id
+
+from .utils import AsyncTyper
+
+_START_WORKER_HELP = "start a datashare worker"
+
+_LIST_WORKFLOWS_HELP = "list registered workflows"
+_LIST_WORKFLOW_NAMES_HELP = "workflow names filters (supports regexes)"
+
+_LIST_ACTIVITIES_HELP = "list registered activities"
+_LIST_ACTIVITY_NAMES_HELP = "activity names filters (supports regexes)"
+
+_START_WORKER_WORKFLOWS_HELP = "workflow names run by the worker (supports regexes)"
+_START_WORKER_ACTIVITIES_HELP = "activity names run by the worker (supports regexes)"
+_START_WORKER_DEPS_HELP = "worker lifetime dependencies name in the registry"
+_START_WORKER_WORKER_ID_PREFIX_HELP = "worker ID prefix"
+_START_WORKER_CONFIG_PATH_HELP = (
+    "path to a worker config YAML file,"
+    " if not provided will load worker configuration from env variables"
+)
+_WORKER_QUEUE_HELP = "worker task queue"
+_TEMPORAL_NAMESPACE_HELP = "worker temporal namespace"
+
+_TEMPORAL_URL_HELP = "address for temporal server"
+_NAMESPACE_HELP = "namespace name"
+_WORKER = "worker"
+
+worker_app = AsyncTyper(name=_WORKER)
+
+logger = logging.getLogger(__name__)
+
+
+@worker_app.async_command(help=_LIST_WORKFLOWS_HELP)
+async def list_workflows(
+    names: Annotated[list[str], typer.Argument(help=_LIST_WORKFLOW_NAMES_HELP)],
+) -> None:
+    workflows = [wf_name for wf_name, _ in discover_workflows(names)]
+    if not workflows:
+        out = """Couldn't find any registered workflow 🤔.
+Make sure your workflow plugins correctly expose workflow entry points, refer to the \
+documentation to learn how to do so."""
+        print(out)
+        return
+    workflows = "\n".join(f"- {wf}" for wf in workflows)
+    out = f"Found {len(workflows)} registered workflows:\n{workflows}"
+    print(out)
+
+
+@worker_app.async_command(help=_LIST_ACTIVITIES_HELP)
+async def list_activities(
+    names: Annotated[list[str], typer.Argument(help=_LIST_ACTIVITY_NAMES_HELP)],
+) -> None:
+    activities = [act_name for act_name, _ in discover_activities(names)]
+    if not activities:
+        out = """Couldn't find any registered activity 🤔.
+    Make sure your activity plugins correctly expose activity entry points, refer \
+to the documentation to learn how to do so."""
+        print(out)
+        return
+    activities = "\n".join(f"- {act}" for act in activities)
+    out = f"Found {len(activities)} registered activities:\n{activities}"
+    print(out)
+
+
+@worker_app.async_command(help=_START_WORKER_HELP)
+async def start(
+    workflows: Annotated[list[str], typer.Option(help=_START_WORKER_WORKFLOWS_HELP)],
+    activities: Annotated[list[str], typer.Option(help=_START_WORKER_ACTIVITIES_HELP)],
+    queue: Annotated[str, typer.Option("--queue", "-q", help=_WORKER_QUEUE_HELP)],
+    dependencies: Annotated[
+        str | None, typer.Option(help=_START_WORKER_DEPS_HELP)
+    ] = None,
+    config_path: Annotated[
+        Path | None,
+        typer.Option(
+            "--config-path", "--config", "-c", help=_START_WORKER_CONFIG_PATH_HELP
+        ),
+    ] = None,
+    worker_id_prefix: Annotated[
+        str | None, typer.Option(help=_START_WORKER_WORKER_ID_PREFIX_HELP)
+    ] = None,
+    temporal_address: Annotated[
+        str | None, typer.Option("--temporal-address", "-a", help=_TEMPORAL_URL_HELP)
+    ] = None,
+    namespace: Annotated[
+        str | None,
+        typer.Option("--temporal-namespace", "-ns", help=_TEMPORAL_NAMESPACE_HELP),
+    ] = None,
+) -> None:
+    if config_path is not None:
+        with config_path.open() as f:
+            bootstrap_config = WorkerConfig.model_validate(
+                yaml.load(f, Loader=yaml.Loader)
+            )
+    else:
+        bootstrap_config = WorkerConfig()
+    temporal_override = dict()
+    if temporal_address is not None:
+        temporal_override["host"] = temporal_address
+    if namespace is not None:
+        temporal_override["namespace"] = namespace
+    if temporal_override:
+        temporal_config = TemporalClientConfig(**temporal_override)
+        update = {"temporal": temporal_config}
+        bootstrap_config = safe_copy(bootstrap_config, update=update)
+    registered_wfs, registered_acts, registered_deps = discover(
+        workflows, act_names=activities, deps_name=dependencies
+    )
+    worker_id = create_worker_id(worker_id_prefix or "worker")
+    client = await bootstrap_config.to_temporal_client()
+    event_loop = asyncio.get_event_loop()
+    async with bootstrap_worker(
+        worker_id,
+        activities=registered_acts,
+        workflows=registered_wfs,
+        dependencies=registered_deps,
+        bootstrap_config=bootstrap_config,
+        client=client,
+        event_loop=event_loop,
+        task_queue=queue,
+    ) as worker:
+        try:
+            await worker.run()
+        except Exception as e:  # noqa: BLE001
+            await worker.shutdown()
+            raise e

{datashare_python-0.2.23 → datashare_python-0.2.25}/datashare_python/config.py

@@ -4,7 +4,13 @@ from icij_common.es import ESClient
 from icij_common.pydantic_utils import ICIJSettings
 from pydantic import Field, PrivateAttr
 from pydantic_settings import SettingsConfigDict
-from temporalio.contrib.pydantic import pydantic_data_converter
+from temporalio.contrib.pydantic import PydanticJSONPlainPayloadConverter, ToJsonOptions
+from temporalio.converter import (
+    CompositePayloadConverter,
+    DataConverter,
+    DefaultPayloadConverter,
+    JSONPlainPayloadConverter,
+)
 
 import datashare_python
 
@@ -64,7 +70,7 @@ class TemporalClientConfig(BaseModel):
             self._client = await TemporalClient.connect(
                 target_host=self.host,
                 namespace=self.namespace,
-                data_converter=pydantic_data_converter,
+                data_converter=PYDANTIC_DATA_CONVERTER,
             )
         return self._client
 
@@ -83,6 +89,8 @@ class WorkerConfig(ICIJSettings, LogWithWorkerIDMixin, BaseModel):
     elasticsearch: ESClientConfig = ESClientConfig()
     temporal: TemporalClientConfig = TemporalClientConfig()
 
+    max_concurrent_io_activities: int = 5
+
     def to_es_client(self) -> ESClient:
         return self.elasticsearch.to_es_client(self.datashare.api_key)
 
@@ -91,3 +99,23 @@ class WorkerConfig(ICIJSettings, LogWithWorkerIDMixin, BaseModel):
 
     async def to_temporal_client(self) -> TemporalClient:
         return await self.temporal.to_client()
+
+
+class _PydanticPayloadConverter(CompositePayloadConverter):
+    def __init__(self) -> None:
+        json_payload_converter = PydanticJSONPlainPayloadConverter(
+            ToJsonOptions(exclude_unset=False)
+        )
+        super().__init__(
+            *(
+                c
+                if not isinstance(c, JSONPlainPayloadConverter)
+                else json_payload_converter
+                for c in DefaultPayloadConverter.default_encoding_payload_converters
+            )
+        )
+
+
+PYDANTIC_DATA_CONVERTER = DataConverter(
+    payload_converter_class=_PydanticPayloadConverter
+)

{datashare_python-0.2.23 → datashare_python-0.2.25}/datashare_python/conftest.py

@@ -1,6 +1,7 @@
 import asyncio
 from asyncio import AbstractEventLoop
 from collections.abc import AsyncGenerator, Generator, Iterator, Sequence
+from pathlib import Path
 
 import aiohttp
 import pytest
@@ -50,6 +51,9 @@ _INDEX_BODY = {
             "language": {"type": "keyword"},
             "documentId": {"type": "keyword"},
             "join": {"type": "join", "relations": {"Document": "NamedEntity"}},
+            "contentType": {"type": "keyword"},
+            "content": {"type": "text"},
+            "contentTranslated": {"type": "text"},
         }
     }
 }
@@ -102,6 +106,13 @@ def test_worker_config() -> WorkerConfig:
     )
 
 
+@pytest.fixture
+def test_worker_config_path(test_worker_config: WorkerConfig, tmpdir: Path) -> Path:
+    config_path = Path(tmpdir) / "config.json"
+    config_path.write_text(test_worker_config.model_dump_json())
+    return config_path
+
+
 @pytest.fixture(scope="session")
 async def worker_lifetime_deps(
     event_loop: AbstractEventLoop,

{datashare_python-0.2.23 → datashare_python-0.2.25}/datashare_python/constants.py

@@ -12,3 +12,5 @@ DEFAULT_TEMPORAL_ADDRESS = "temporal:7233"
 DEFAULT_DS_ADDRESS = "http://localhost:8080"
 
 DEFAULT_NAMESPACE = "datashare-default"
+
+METADATA_JSON = "metadata.json"

{datashare_python-0.2.23 → datashare_python-0.2.25}/datashare_python/dependencies.py

@@ -1,8 +1,11 @@
+import inspect
 import logging
 from asyncio import AbstractEventLoop, iscoroutine
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Callable
 from contextlib import AsyncExitStack, asynccontextmanager
 from contextvars import ContextVar
+from copy import deepcopy
+from typing import Any
 
 from icij_common.es import ESClient
 
@@ -20,7 +23,7 @@ TASK_CLIENT: ContextVar[DatashareTaskClient] = ContextVar("task_client")
 TEMPORAL_CLIENT: ContextVar[TemporalClient] = ContextVar("temporal_client")
 
 
-def set_event_loop(event_loop: AbstractEventLoop, **_) -> None:
+def set_event_loop(event_loop: AbstractEventLoop) -> None:
     EVENT_LOOP.set(event_loop)
 
 
@@ -31,13 +34,13 @@ def lifespan_event_loop() -> AbstractEventLoop:
         raise DependencyInjectionError("event loop") from e
 
 
-def set_loggers(worker_config: WorkerConfig, worker_id: str, **_) -> None:
+def set_loggers(worker_config: WorkerConfig, worker_id: str) -> None:
     worker_config.setup_loggers(worker_id=worker_id)
     logger.info("worker loggers ready to log 💬")
     logger.info("app config: %s", worker_config.model_dump_json(indent=2))
 
 
-async def set_es_client(worker_config: WorkerConfig, **_) -> ESClient:
+async def set_es_client(worker_config: WorkerConfig) -> ESClient:
     client = worker_config.to_es_client()
     ES_CLIENT.set(client)
     return client
@@ -52,7 +55,7 @@ def lifespan_es_client() -> ESClient:
 
 
 # Task client setup
-async def set_task_client(worker_config: WorkerConfig, **_) -> DatashareTaskClient:
+async def set_task_client(worker_config: WorkerConfig) -> DatashareTaskClient:
     task_client = worker_config.to_task_client()
     TASK_CLIENT.set(task_client)
     return task_client
@@ -67,7 +70,7 @@ def lifespan_task_client() -> DatashareTaskClient:
 
 
 # Temporal client setup
-async def set_temporal_client(worker_config: WorkerConfig, **_) -> None:
+async def set_temporal_client(worker_config: WorkerConfig) -> None:
     client = await worker_config.to_temporal_client()
     TEMPORAL_CLIENT.set(client)
 
@@ -86,7 +89,7 @@ async def with_dependencies(
 ) -> AsyncGenerator[None, None]:
     async with AsyncExitStack() as stack:
         for dep in dependencies:
-            cm = dep(**kwargs)
+            cm = dep(**add_missing_args(dep, dict(), **kwargs))
             if hasattr(cm, "__aenter__"):
                 await stack.enter_async_context(cm)
             elif hasattr(cm, "__enter__"):
@@ -94,3 +97,20 @@ async def with_dependencies(
             elif iscoroutine(cm):
                 await cm
         yield
+
+
+def add_missing_args(fn: Callable, args: dict[str, Any], **kwargs) -> dict[str, Any]:
+    # We make the choice not to raise in case of missing argument here, the error will
+    # be correctly raise when the function is called
+    from_kwargs = dict()
+    sig = inspect.signature(fn)
+    for param_name in sig.parameters:
+        if param_name in args:
+            continue
+        kwargs_value = kwargs.get(param_name)
+        if kwargs_value is not None:
+            from_kwargs[param_name] = kwargs_value
+    if from_kwargs:
+        args = deepcopy(args)
+        args.update(from_kwargs)
+    return args

datashare_python-0.2.25/datashare_python/discovery.py

@@ -0,0 +1,178 @@
+import logging
+import re
+from collections.abc import Callable, Iterable
+from importlib.metadata import entry_points
+
+from .types_ import ContextManagerFactory
+from .utils import ActivityWithProgress
+
+logger = logging.getLogger(__name__)
+
+Activity = ActivityWithProgress | Callable | type
+
+_DEPENDENCIES = "dependencies"
+_WORKFLOW_GROUPS = "datashare.workflows"
+_ACTIVITIES_GROUPS = "datashare.activities"
+_DEPENDENCIES_GROUPS = "datashare.dependencies"
+
+_RegisteredWorkflow = tuple[str, type]
+_RegisteredActivity = tuple[str, Activity]
+_Dependencies = list[ContextManagerFactory]
+_Discovery = tuple[
+    Iterable[_RegisteredWorkflow] | None,
+    Iterable[_RegisteredActivity] | None,
+    _Dependencies | None,
+]
+
+
+def discover(
+    wf_names: list[str] | None, *, act_names: list[str] | None, deps_name: str | None
+) -> _Discovery:
+    discovered = ""
+    wfs = None
+    if wf_names is not None:
+        discovered_wfs = discover_workflows(wf_names)
+        if discovered_wfs:
+            wf_names, wfs = zip(*discovered_wfs, strict=True)
+            if wf_names:
+                n_wfs = len(wf_names)
+                discovered += (
+                    f"- {n_wfs} workflow{'s' if n_wfs > 1 else ''}:"
+                    f" {', '.join(wf_names)}"
+                )
+    acts = None
+    if act_names is not None:
+        discovered_acts = discover_activities(act_names)
+        if discovered_acts:
+            act_names, acts = zip(*discovered_acts, strict=True)
+            if act_names:
+                if discovered:
+                    discovered += "\n"
+                n_acts = len(act_names)
+                discovered += (
+                    f"- {n_acts} activit{'ies' if n_acts > 1 else 'y'}:"
+                    f" {', '.join(act_names)}"
+                )
+    if not acts and not wfs:
+        raise ValueError("Couldn't find any registered activity or workflow.")
+    deps = discover_dependencies(deps_name)
+    if deps:
+        n_deps = len(deps)
+        discovered += "\n"
+        deps_names = (d.__name__ for d in deps)
+        discovered += (
+            f"- {n_deps} dependenc{'ies' if n_deps > 1 else 'y'}:"
+            f" {', '.join(deps_names)}"
+        )
+    logger.info("discovered:\n%s", discovered)
+    return wfs, acts, deps
+
+
+def discover_workflows(names: list[str]) -> list[_RegisteredWorkflow]:
+    pattern = None if not names else re.compile(rf"^{'|'.join(names)}$")
+    impls = entry_points(group=_WORKFLOW_GROUPS)
+    registered = []
+    for wf_impls in impls:
+        wf_impls = wf_impls.load()  # noqa: PLW2901
+        if not isinstance(wf_impls, list | tuple | set):
+            wf_impls = [wf_impls]  # noqa: PLW2901
+        for wf_impl in wf_impls:
+            wf_name = _parse_wf_name(wf_impl)
+            if pattern and not pattern.match(wf_name):
+                continue
+            registered.append((wf_name, wf_impl))
+    return registered
+
+
+def discover_activities(names: list[str]) -> list[_RegisteredActivity]:
+    pattern = None if not names else re.compile(rf"^{'|'.join(names)}$")
+    impls = entry_points(group=_ACTIVITIES_GROUPS)
+    registered = []
+    for act_impls in impls:
+        act_impls = act_impls.load()  # noqa: PLW2901
+        if not isinstance(act_impls, list | tuple | set):
+            act_impls = [act_impls]  # noqa: PLW2901
+        for act_impl in act_impls:
+            act_name = _parse_activity_name(act_impl)
+            if pattern and not pattern.match(act_name):
+                continue
+            registered.append((act_name, act_impl))
+    return registered
+
+
+def discover_dependencies(name: str | None) -> _Dependencies | None:
+    impls = entry_points(name=_DEPENDENCIES, group=_DEPENDENCIES_GROUPS)
+    if not impls:
+        if name is None:
+            return None
+        available_impls = entry_points(group=_DEPENDENCIES_GROUPS)
+        msg = (
+            f'failed to find dependency: "{name}", '
+            f"available dependencies: {available_impls}"
+        )
+        raise LookupError(msg)
+    if len(impls) > 1:
+        msg = f'found multiple dependencies for name "{name}": {impls}'
+        raise ValueError(msg)
+    deps_registry = impls[_DEPENDENCIES].load()
+    if name:
+        try:
+            return deps_registry[name]
+        except KeyError as e:
+            available = list(deps_registry)
+            msg = (
+                f'failed to find dependency for name "{name}", available dependencies: '
+                f"{available}"
+            )
+            raise LookupError(msg) from e
+    if not deps_registry:
+        raise ValueError("empty dependency registry !")
+    if len(deps_registry) > 1:
+        available = ", ".join('"' + d + '"' for d in deps_registry)
+        msg = (
+            f"dependency registry contains multiples entries {available},"
+            f" please select one by providing a name"
+        )
+        raise ValueError(msg)
+    return next(iter(deps_registry.values()))
+
+
+def _parse_wf_name(wf_type: type) -> str:
+    if not isinstance(wf_type, type):
+        msg = (
+            f"expected registered workflow implementation to be a temporal workflow"
+            f" decorated with @workflow.defn(name=<name>) class, found: {type(wf_type)}"
+        )
+        raise TypeError(msg)
+
+    wf_defn = getattr(wf_type, "__temporal_workflow_definition", None)
+    if wf_defn is None:
+        msg = (
+            f"expected registered workflow implementation to be a temporal workflow"
+            f" decorated with @workflow.defn(name=<name>) class, found: {wf_type}"
+        )
+        raise ValueError(msg)
+    if wf_defn.name is None:
+        msg = (
+            "missing workflow definition name, please register your workflow"
+            " with an explicit name: @workflow.defn(name=<name>)"
+        )
+        raise ValueError(msg)
+    return wf_defn.name
+
+
+def _parse_activity_name(act: Activity) -> str:
+    act_defn = getattr(act, "__temporal_activity_definition", None)
+    if act_defn is None:
+        msg = (
+            f"expected registered actitiby implementation to be a temporal activity"
+            f" decorated with @activity.defn(name=<name>), found: {act}"
+        )
+        raise ValueError(msg)
+    if act_defn.name is None:
+        msg = (
+            "missing activity definition name, please register your activities"
+            " with an explicit name: @activity.defn(name=<name>)"
+        )
+        raise ValueError(msg)
+    return act_defn.name

{datashare_python-0.2.23 → datashare_python-0.2.25}/datashare_python/objects.py

@@ -1,9 +1,9 @@
 import logging
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from enum import StrEnum, unique
-from typing import Any, Literal, Self
+from typing import Any, Literal, Self, TypeVar
 
 from temporalio import workflow
 
@@ -23,14 +23,26 @@ from pydantic.main import IncEx
 logger = logging.getLogger(__name__)
 
 
+T = TypeVar("T")
+Predicate = Callable[[T], bool] | Callable[[T], Awaitable[bool]]
+
+
 class BaseModel(_BaseModel):
     model_config = merge_configs(icij_config(), no_enum_values_config())
 
 
+class BasePayload(_BaseModel):
+    model_config = icij_config()
+
+
 class DatashareModel(BaseModel):
     model_config = merge_configs(BaseModel.model_config, lowercamel_case_config())
 
 
+class LowerCamelCaseModel(_BaseModel):
+    model_config = merge_configs(icij_config(), lowercamel_case_config())
+
+
 @unique
 class TaskState(StrEnum):
     CREATED = "CREATED"
@@ -135,12 +147,14 @@ class TaskGroup:
 class Document(DatashareModel):
     id: str
     root_document: str
-    content: str
     language: str
+    content: str | None = None
+    content_type: str | None = None
     tags: list[str] = Field(default_factory=list)
     content_translated: dict[str, str] = Field(
         default_factory=dict, alias="content_translated"
     )
+    type: str = Field(default="Document", frozen=True)
 
     @classmethod
     def from_es(cls, es_doc: dict) -> Self:

{datashare_python-0.2.23 → datashare_python-0.2.25}/datashare_python/template.py

@@ -71,12 +71,12 @@ def _update_pyproject_toml(
 ) -> dict[str, Any]:
     pyproject_toml = deepcopy(pyproject_toml)
 
-    pyproject_toml["tool"]["uv"].pop("sources")
-    pyproject_toml["tool"]["uv"].pop("index")
+    pyproject_toml["tool"]["uv"].pop("sources", None)
+    pyproject_toml["tool"]["uv"].pop("index", None)
 
     project = pyproject_toml["project"]
     project["authors"] = []
-    project.pop("urls")
+    project.pop("urls", None)
     project["dependencies"] = sorted(
         d
         for d in project["dependencies"]
@@ -87,7 +87,7 @@ def _update_pyproject_toml(
         for d in project["dependencies"]
         if any(d.startswith(base) for base in _BASE_DEPS)
     )
-    project.pop("optional-dependencies")
+    project.pop("optional-dependencies", None)
 
     entry_points = project["entry-points"]
 
@@ -105,4 +105,11 @@ def _update_pyproject_toml(
         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"]
+        ]
+
     return pyproject_toml