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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datashare-python
-Version: 0.8.20
+Version: 0.9.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/
@@ -13,6 +13,7 @@ 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: lru-dict~=1.4
 Requires-Dist: nest-asyncio~=1.6
 Requires-Dist: orjson~=3.11
 Requires-Dist: pydantic-extra-types[pycountry]>=2.11.1

{datashare_python-0.8.20 → datashare_python-0.9.0}/datashare_python/config.py

@@ -6,13 +6,6 @@ from icij_common.es import ESClient
 from icij_common.pydantic_utils import ICIJSettings
 from pydantic import PrivateAttr
 from pydantic_settings import SettingsConfigDict
-from temporalio.contrib.pydantic import PydanticJSONPlainPayloadConverter, ToJsonOptions
-from temporalio.converter import (
-    CompositePayloadConverter,
-    DataConverter,
-    DefaultPayloadConverter,
-    JSONPlainPayloadConverter,
-)
 from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
 
 import datashare_python
@@ -20,6 +13,7 @@ import datashare_python
 from .objects import BaseModel
 from .task_client import DatashareTaskClient
 from .types_ import TemporalClient
+from .utils import PYDANTIC_DATA_CONVERTER
 
 _ALL_LOGGERS = [datashare_python.__name__]
 
@@ -129,23 +123,3 @@ class WorkerConfig(ICIJSettings, 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.8.20 → datashare_python-0.9.0}/datashare_python/constants.py

@@ -13,3 +13,5 @@ DEFAULT_NAMESPACE = "datashare-default"
 METADATA_JSON = "metadata.json"
 
 TIKA_METADATA_RESOURCENAME = "tika_metadata_resourcename"
+
+DEFAULT_SHARED_RESOURCES_SIZE = 1

{datashare_python-0.8.20 → datashare_python-0.9.0}/datashare_python/dependencies.py

@@ -2,7 +2,7 @@ import inspect
 import logging
 from asyncio import AbstractEventLoop, iscoroutine
 from collections.abc import AsyncGenerator, Callable
-from contextlib import AsyncExitStack, asynccontextmanager
+from contextlib import AbstractContextManager, AsyncExitStack, asynccontextmanager
 from contextvars import ContextVar
 from copy import deepcopy
 from typing import Any
@@ -12,6 +12,7 @@ from icij_common.es import ESClient
 from .config import LogLevel, WorkerConfig
 from .exceptions import DependencyInjectionError
 from .logging_ import setup_worker_loggers
+from .objects import Shared
 from .task_client import DatashareTaskClient
 from .types_ import ContextManagerFactory, TemporalClient
 
@@ -23,6 +24,7 @@ ES_CLIENT: ContextVar[ESClient] = ContextVar("es_client")
 TASK_CLIENT: ContextVar[DatashareTaskClient] = ContextVar("task_client")
 TEMPORAL_CLIENT: ContextVar[TemporalClient] = ContextVar("temporal_client")
 WORKER_CONFIG: ContextVar[WorkerConfig] = ContextVar("worker_config")
+SHARED: ContextVar[Shared] = ContextVar("shared")
 
 
 def set_event_loop(event_loop: AbstractEventLoop) -> None:
@@ -99,6 +101,20 @@ def lifespan_temporal_client() -> TemporalClient:
         raise DependencyInjectionError("temporal client") from e
 
 
+# Setup shared resources
+async def set_shared_resources(shared: Shared) -> Shared:
+    SHARED.set(shared)
+    return shared
+
+
+# Return shared resources
+def lifespan_shared_resources() -> Shared:
+    try:
+        return SHARED.get()
+    except LookupError as e:
+        raise DependencyInjectionError("shared resources") from e
+
+
 @asynccontextmanager
 async def with_dependencies(
     dependencies: list[ContextManagerFactory], **kwargs
@@ -130,3 +146,11 @@ def add_missing_args(fn: Callable, args: dict[str, Any], **kwargs) -> dict[str,
         args = deepcopy(args)
         args.update(from_kwargs)
     return args
+
+
+# component lifecycle
+def component_teardown(_cache_key: str, component: AbstractContextManager) -> None:
+    if not isinstance(component, AbstractContextManager):
+        return
+
+    component.__exit__(None, None, None)

{datashare_python-0.8.20 → datashare_python-0.9.0}/datashare_python/interceptors.py

@@ -1,13 +1,31 @@
+import asyncio
+import contextlib
+import dataclasses
 import secrets
-from collections.abc import Generator, Mapping
+from collections.abc import Callable, Generator, Mapping
 from contextlib import contextmanager
 from contextvars import ContextVar
 from copy import deepcopy
-from typing import Annotated, Any, NoReturn, Self, TypeVar
+from functools import partial
+from inspect import signature
+from types import UnionType
+from typing import (
+    Annotated,
+    Any,
+    NoReturn,
+    Self,
+    TypeVar,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
 
 from nexusrpc import InputT, OutputT
 from pydantic import Field
+from temporalio import activity
+from temporalio.activity import _Definition
 from temporalio.api.common.v1 import Payload
+from temporalio.client import WorkflowHandle
 from temporalio.converter import DataConverter
 from temporalio.worker import (
     ActivityInboundInterceptor,
@@ -34,6 +52,13 @@ from temporalio.workflow import (
 )
 
 from .objects import BaseModel
+from .types_ import ProgressRateHandler, Weight
+from .utils import (
+    PROGRESS_HANDLER_ARG,
+    PYDANTIC_DATA_CONVERTER,
+    ActivityWithProgress,
+    ProgressSignal,
+)
 
 _TRACEPARENT = "traceparent"
 _DEFAULT_PAYLOAD_CONVERTER = DataConverter.default.payload_converter
@@ -196,3 +221,156 @@ def _with_trace_context_header(
         next_ctx.traceparent
     )
     return new_obj
+
+
+class ProgressInterceptor(Interceptor):
+    def intercept_activity(
+        self,
+        next: ActivityInboundInterceptor,  # noqa: A002
+    ) -> ActivityInboundInterceptor:
+        return _ProgressInboundInterceptor(next)
+
+
+def _parse_progress_weight(act_fn: Callable) -> float:
+    hints = get_type_hints(act_fn, include_extras=True)
+    hint = hints["progress"]
+    annotated_progress = get_origin(hint)
+    if annotated_progress is not Annotated:
+        return 1.0
+    annotated_args = get_args(hint)
+    for ann in annotated_args[1:]:
+        if isinstance(ann, Weight):
+            return ann.value
+    return 1.0
+
+
+async def progress_handler(
+    progress: float,
+    handle: WorkflowHandle,
+    *,
+    activity_id: str,
+    run_id: str,
+    weight: float = 1.0,
+) -> None:
+    signal = ProgressSignal(
+        activity_id=activity_id, run_id=run_id, progress=progress, weight=weight
+    )
+    await handle.signal("update_progress", signal)
+
+
+def supports_progress(task_fn: Callable) -> bool:
+    return any(
+        param.name == PROGRESS_HANDLER_ARG
+        for param in signature(task_fn).parameters.values()
+    )
+
+
+def _get_progress_handler(act_fn: Callable) -> ProgressRateHandler:
+    act = getattr(act_fn, "__self__", None)
+    # Weirdly isinstance doesn't work here
+    if act is None or not isinstance(act, ActivityWithProgress):
+        msg = (
+            f"to support progress, activities should inherit from "
+            f"{ActivityWithProgress.__name__}."
+        )
+        raise TypeError(msg)
+    weight = _parse_progress_weight(act_fn)
+    info = activity.info()
+    run_id = info.workflow_run_id
+    workflow_id = info.workflow_id
+    activity_id = activity.info().activity_id
+    client = act._temporal_client
+    workflow_handle = client.get_workflow_handle(workflow_id, run_id=run_id)
+    handler = partial(
+        progress_handler,
+        handle=workflow_handle,
+        run_id=run_id,
+        activity_id=activity_id,
+        weight=weight,
+    )
+    return handler
+
+
+def _is_progress(t: type) -> bool:
+    if t is ProgressRateHandler:
+        return True
+    return bool(
+        isinstance(t, UnionType)
+        and any(sub_t is ProgressRateHandler for sub_t in get_args(t))
+    )
+
+
+def _without_progress(arg_types: list[type] | None) -> list[type] | None:
+    if arg_types is None:
+        return None
+    filtered = [t for t in arg_types if not _is_progress(t)]
+    return filtered
+
+
+class _ProgressInboundInterceptor(ActivityInboundInterceptor):
+    async def execute_activity(self, input: ExecuteActivityInput) -> Any:  # noqa: A002
+        if not supports_progress(input.fn):
+            return await super().execute_activity(input)
+        # The progress args breaks trigger a bypass of the dataloader:
+        # https://github.com/temporalio/sdk-python/blob/631ebaf0e20fb214b16589b45627b358048a5d77/temporalio/worker/_activity.py#L600
+        # we have to force it here again
+        progress_handler = _get_progress_handler(input.fn)
+        if input.args:
+            data_converter = PYDANTIC_DATA_CONVERTER
+            arg_types = _Definition.must_from_callable(input.fn).arg_types
+            arg_types = _without_progress(arg_types)
+            arg_types = arg_types[: len(input.args)]
+            encoded = await data_converter.encode(input.args)
+            new_args = await data_converter.decode(encoded, type_hints=arg_types)
+            new_args.append(progress_handler)
+        else:
+            new_args = [progress_handler]
+        new_input = dataclasses.replace(input, args=new_args)
+        await progress_handler(0.0)
+        res = await super().execute_activity(new_input)
+        await progress_handler(1.0)
+        return res
+
+
+class HeartbeatInterceptor(Interceptor):
+    def __init__(self, n_missed_before_timeout: int = 5):
+        self._n_missed_before_timeout = n_missed_before_timeout
+
+    def intercept_activity(
+        self,
+        next: ActivityInboundInterceptor,  # noqa: A002
+    ) -> ActivityInboundInterceptor:
+        return _HeartbeatInboundInterceptor(next, self._n_missed_before_timeout)
+
+
+async def _heartbeat_every(period: float, *details: Any) -> None:
+    with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
+        activity.heartbeat(*details)
+    while True:
+        await asyncio.sleep(period)
+        with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
+            activity.heartbeat(*details)
+
+
+class _HeartbeatInboundInterceptor(ActivityInboundInterceptor):
+    def __init__(
+        self,
+        next: ActivityInboundInterceptor,  # noqa: A002
+        n_missed_before_timeout: int = 5,
+    ) -> None:
+        super().__init__(next)
+        self._n_missed_before_timeout = n_missed_before_timeout
+
+    async def execute_activity(self, input: ExecuteActivityInput) -> Any:  # noqa: A002
+        heartbeat_timeout = activity.info().heartbeat_timeout
+        heartbeat_task = None
+        if heartbeat_timeout:
+            period = heartbeat_timeout.total_seconds() / self._n_missed_before_timeout
+            heartbeat_task = asyncio.create_task(_heartbeat_every(period))
+        try:
+            activity.heartbeat()
+            return await super().execute_activity(input)
+        finally:
+            if heartbeat_task:
+                heartbeat_task.cancel()
+                await asyncio.wait([heartbeat_task])

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

@@ -1,7 +1,8 @@
 import logging
 import os
+from asyncio import Lock
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass
+from dataclasses import InitVar, dataclass, field
 from datetime import UTC, datetime
 from enum import StrEnum, unique
 from io import BytesIO
@@ -9,12 +10,13 @@ from pathlib import Path
 from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, cast
 
 import langcodes
+from lru import LRU
 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
+from .constants import DEFAULT_SHARED_RESOURCES_SIZE, TIKA_METADATA_RESOURCENAME
 
 with workflow.unsafe.imports_passed_through():
     from icij_common.es import (
@@ -367,3 +369,47 @@ 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)

{datashare_python-0.8.20 → datashare_python-0.9.0}/datashare_python/types_.py

@@ -1,5 +1,6 @@
 from collections.abc import Coroutine
 from contextlib import AbstractAsyncContextManager, AbstractContextManager
+from dataclasses import dataclass
 from typing import Protocol
 
 from temporalio.client import Client
@@ -12,6 +13,11 @@ class ProgressRateHandler(Protocol):
         pass
 
 
+@dataclass
+class Weight:
+    value: float
+
+
 class RawProgressHandler(Protocol):
     async def __call__(self, iteration: int) -> None:
         pass

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

@@ -1,18 +1,15 @@
 import asyncio
 import contextlib
-import contextvars
 import inspect
 import json
 import os
 import shutil
-import threading
-from collections.abc import Awaitable, Callable, Coroutine, Iterable
+from collections.abc import Callable, Coroutine, Iterable, Sequence
 from copy import deepcopy
 from dataclasses import dataclass
 from datetime import timedelta
-from functools import partial, wraps
+from functools import wraps
 from hashlib import sha256
-from inspect import signature
 from io import BytesIO
 from pathlib import Path
 from typing import Any, ParamSpec, TypeVar
@@ -20,9 +17,18 @@ from uuid import uuid4
 
 import nest_asyncio
 import temporalio
+from pydantic import ValidationError
 from temporalio import activity, workflow
-from temporalio.client import Client, WorkflowHandle
+from temporalio.api.common.v1 import Payload
+from temporalio.client import Client
 from temporalio.common import RetryPolicy, SearchAttributeKey
+from temporalio.contrib.pydantic import PydanticJSONPlainPayloadConverter, ToJsonOptions
+from temporalio.converter import (
+    CompositePayloadConverter,
+    DataConverter,
+    DefaultPayloadConverter,
+    JSONPlainPayloadConverter,
+)
 from temporalio.exceptions import ApplicationError
 
 from .constants import METADATA_JSON
@@ -98,9 +104,14 @@ class WorkflowWithProgress:
 
 def _retry_policy_with_default(retry_policy: RetryPolicy | None) -> RetryPolicy:
     if retry_policy is None:
-        retry_policy = RetryPolicy(non_retryable_error_types=[])
+        retry_policy = RetryPolicy(non_retryable_error_types=[], maximum_attempts=3)
     retry_policy = deepcopy(retry_policy)
-    non_retryable_error_types = set(retry_policy.non_retryable_error_types)
+    non_retryable_error_types = (
+        retry_policy.non_retryable_error_types
+        if retry_policy.non_retryable_error_types is not None
+        else []
+    )
+    non_retryable_error_types = set(non_retryable_error_types)
     non_retryable_error_types.update(_NEVER_RETRIABLES)
     retry_policy.non_retryable_error_types = list(non_retryable_error_types)
     return retry_policy
@@ -130,167 +141,6 @@ async def execute_activity(
     )
 
 
-async def progress_handler(
-    progress: float,
-    handle: WorkflowHandle,
-    *,
-    activity_id: str,
-    run_id: str,
-    weight: float = 1.0,
-) -> None:
-    signal = ProgressSignal(
-        activity_id=activity_id, run_id=run_id, progress=progress, weight=weight
-    )
-    await handle.signal("update_progress", signal)
-    with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
-        activity.heartbeat()
-
-
-def get_activity_progress_handler_async(
-    client: Client, weight: float
-) -> ProgressRateHandler:
-    info = activity.info()
-    run_id = info.workflow_run_id
-    workflow_id = info.workflow_id
-    activity_id = activity.info().activity_id
-    workflow_handle = client.get_workflow_handle(workflow_id, run_id=run_id)
-    handler = partial(
-        progress_handler,
-        handle=workflow_handle,
-        run_id=run_id,
-        activity_id=activity_id,
-        weight=weight,
-    )
-    return handler
-
-
-def supports_progress(task_fn: Callable) -> bool:
-    return any(
-        param.name == PROGRESS_HANDLER_ARG
-        for param in signature(task_fn).parameters.values()
-    )
-
-
-def with_progress(weight: float = 1.0) -> Callable[P, T]:
-    if isinstance(weight, Callable):
-        return with_progress(weight=1)(weight)
-
-    def decorator(activity_fn: Callable[P, T]) -> Callable[P, T]:
-        # TODO: handle the fact activities should have only positional args...
-        if asyncio.iscoroutinefunction(activity_fn):
-
-            @wraps(activity_fn)
-            async def wrapper(self: ActivityWithProgress, *args: P.args) -> T:
-                if not isinstance(self, ActivityWithProgress):
-                    msg = (
-                        f"{with_progress.__name__} decorator is meant to be used on "
-                        f"activities defined as an {ActivityWithProgress.__name__}"
-                        f" method, expected a {ActivityWithProgress.__name__} as first"
-                        f" argument, found {self}"
-                    )
-                    raise TypeError(msg)
-                handler = get_activity_progress_handler_async(
-                    client=self._temporal_client, weight=weight
-                )
-                await handler(0.0)
-                res = await activity_fn(self, *args, progress=handler)
-                await handler(1.0)
-                return res
-
-        else:
-
-            @wraps(activity_fn)
-            def wrapper(self: ActivityWithProgress, *args: P.args) -> T:
-                if not isinstance(self, ActivityWithProgress):
-                    msg = (
-                        f"{with_progress.__name__} decorator is meant to be used on "
-                        f"activities defined as an {ActivityWithProgress.__name__}"
-                        f" method, expected a {ActivityWithProgress.__name__} as first"
-                        f" argument, found {self}"
-                    )
-                    raise TypeError(msg)
-                handler = get_activity_progress_handler_async(
-                    client=self._temporal_client, weight=weight
-                )
-                event_loop = self._event_loop
-                asyncio.run_coroutine_threadsafe(handler(0.0), event_loop).result()
-                res = activity_fn(self, *args, progress=handler)
-                asyncio.run_coroutine_threadsafe(handler(1.0), event_loop).result()
-                return res
-
-        return wrapper
-
-    return decorator
-
-
-def with_async_heartbeat(
-    activity_fn: Callable[P, Awaitable[T]], n_missed_before_timeout: int
-) -> Callable[P, Awaitable[T]]:
-    # Copied from
-    # https://github.com/temporalio/samples-python/blob/main/custom_decorator/activity_utils.py
-    @wraps(activity_fn)
-    async def wrapper(*args, **kwargs) -> T:
-        heartbeat_timeout = activity.info().heartbeat_timeout
-        heartbeat_task = None
-        if heartbeat_timeout:
-            period = heartbeat_timeout.total_seconds() / n_missed_before_timeout
-            heartbeat_task = asyncio.create_task(_async_heartbeat_every(period))
-        try:
-            activity.heartbeat()
-            return await activity_fn(*args, **kwargs)
-        finally:
-            if heartbeat_task:
-                heartbeat_task.cancel()
-                await asyncio.wait([heartbeat_task])
-
-    return wrapper
-
-
-async def _async_heartbeat_every(period: float, *details: Any) -> None:
-    with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
-        activity.heartbeat(*details)
-    while True:
-        await asyncio.sleep(period)
-        with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
-            activity.heartbeat(*details)
-
-
-def with_sync_heartbeat(
-    activity_fn: Callable[P, T], n_missed_before_timeout: int
-) -> Callable[P, T]:
-    @wraps(activity_fn)
-    def wrapper(*args, **kwargs) -> T:
-        heartbeat_timeout = activity.info().heartbeat_timeout
-        heartbeat_thread, stop_event = None, None
-        if heartbeat_timeout:
-            period = heartbeat_timeout.total_seconds() / n_missed_before_timeout
-            ctx = contextvars.copy_context()
-            run_args = (_sync_heartbeat_every, period, threading.Event())
-            heartbeat_thread, stop_event = (
-                threading.Thread(target=ctx.run, args=run_args),
-                run_args[-1],
-            )
-            heartbeat_thread.start()
-        try:
-            return activity_fn(*args, **kwargs)
-        finally:
-            if heartbeat_thread:
-                stop_event.set()
-                heartbeat_thread.join()
-
-    return wrapper
-
-
-def _sync_heartbeat_every(
-    period: float, stop_event: threading.Event, *details: Any
-) -> None:
-    with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
-        activity.heartbeat(*details)
-    while not stop_event.wait(period):
-        with contextlib.suppress(RuntimeError, asyncio.TimeoutError):
-            activity.heartbeat(*details)
-
-
 def positional_args_only(activity_fn: Callable[P, T]) -> Callable[P, T]:
     sig = inspect.signature(activity_fn)
 
@@ -395,29 +245,27 @@ def with_retriables(
 
 
 def activity_defn(
-    name: str,
-    progress_weight: float = 1.0,
-    retriables: set[type[Exception]] = None,
-    n_missed_heartbeats_before_timeout: int = 5,
+    name: str, retriables: set[type[Exception]] = None
 ) -> Callable[[Callable[P, T]], Callable[P, T]]:
+
     def decorator(activity_fn: Callable[P, T]) -> Callable[P, T]:
-        # TODO: some of these could probably be reimplemented more elegantly using
-        #  temporal interceptors: https://docs.temporal.io/develop/python/workers/interceptors
         activity_fn = positional_args_only(activity_fn)
         activity_fn = with_retriables(retriables)(activity_fn)
-        if supports_progress(activity_fn):
-            activity_fn = with_progress(progress_weight)(activity_fn)
+        activity_fn = activity.defn(activity_fn, name=name)
+
         is_async = asyncio.iscoroutinefunction(activity_fn)
         if is_async:
-            activity_fn = with_async_heartbeat(
-                activity_fn, n_missed_heartbeats_before_timeout
-            )
+
+            @wraps(activity_fn)
+            async def wrapper(*args, **kwargs) -> T:
+                return await activity_fn(*args, **kwargs)
         else:
-            activity_fn = with_sync_heartbeat(
-                activity_fn, n_missed_heartbeats_before_timeout
-            )
-        activity_fn = activity.defn(activity_fn, name=name)
-        return activity_fn
+
+            @wraps(activity_fn)
+            def wrapper(*args, **kwargs) -> T:
+                return activity_fn(*args, **kwargs)
+
+        return wrapper
 
     return decorator
 
@@ -488,6 +336,7 @@ def write_artifact(root: Path, artifact: DocArtifact) -> Path:
             with artifact_path.open("wb") as f:
                 f.write(artifact.artifact.read())
         case Path():
+            artifact_path.unlink(missing_ok=True)
             shutil.move(artifact.artifact, artifact_path)
         case _:
             msg = f"unsupported artifact type: {artifact.artifact.__class__.__name__}"
@@ -601,3 +450,37 @@ def read_jsonl(path: Path) -> Iterable[dict]:
             line = line.strip()  # noqa: PLW2901
             if line:
                 yield json.loads(line)
+
+
+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
+            )
+        )
+
+    def from_payloads(
+        self, payloads: Sequence[Payload], type_hints: list[type] | None = None
+    ) -> list[Any]:
+        try:
+            return super().from_payloads(payloads, type_hints)
+        except (TypeError, ValidationError) as e:
+            raise fatal_error_from_exception(e) from e
+
+    def to_payloads(self, values: Sequence[Any]) -> list[Payload]:
+        try:
+            return super().to_payloads(values)
+        except (TypeError, ValidationError) as e:
+            raise fatal_error_from_exception(e) from e
+
+
+PYDANTIC_DATA_CONVERTER = DataConverter(
+    payload_converter_class=_PydanticPayloadConverter
+)

{datashare_python-0.8.20 → datashare_python-0.9.0}/datashare_python/worker.py

@@ -21,7 +21,11 @@ from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
 from .config import WorkerConfig
 from .dependencies import with_dependencies
 from .discovery import Activity
-from .interceptors import TraceContextInterceptor
+from .interceptors import (
+    HeartbeatInterceptor,
+    ProgressInterceptor,
+    TraceContextInterceptor,
+)
 from .types_ import ContextManagerFactory, TemporalClient
 
 logger = logging.getLogger(__name__)
@@ -91,7 +95,11 @@ def datashare_worker(
         max_concurrent_activities = 1
         if workflows:
             logger.warning(_SEPARATE_IO_AND_CPU_WORKERS)
-    interceptors = [TraceContextInterceptor()]
+    interceptors = [
+        TraceContextInterceptor(),
+        ProgressInterceptor(),
+        HeartbeatInterceptor(),
+    ]
     wf_runner = SandboxedWorkflowRunner() if sandboxed else UnsandboxedWorkflowRunner()
     return DatashareWorker(
         client,

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

@@ -1,6 +1,6 @@
 [project]
 name = "datashare-python"
-version = "0.8.20"
+version = "0.9.0"
 description = "Manage Python tasks and local resources in Datashare"
 authors = [
   { name = "Clément Doumouro", email = "cdoumouro@icij.org" },
@@ -23,6 +23,7 @@ dependencies = [
   "temporalio~=1.23",
   "typer>=0.15.4,<0.25.1",
   "tomlkit~=0.14.0",
+  "lru-dict~=1.4",
 ]
 
 [project.urls]