Allow nested dataclasses inside lists, tuples and dicts as task args (#5)

lukasbindreiter · web-flow · commit 978606c40486 · 2025-06-24T12:45:32.000+02:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -6,7 +6,7 @@ repos:
       - id: end-of-file-fixer
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     # keep the version here in sync with the version in uv.lock
-    rev: "v0.11.13"
+    rev: "v0.12.0"
     hooks:
       - id: ruff-check
         args: [--fix, --exit-non-zero-on-fix]
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -17,7 +17,7 @@
   "python.analysis.autoImportCompletions": true,
   "python.analysis.typeCheckingMode": "standard",
   "python.testing.pytestArgs": [
-    "tilebox-datasets"
+    "tilebox-workflows"
   ],
   "search.useGlobalIgnoreFiles": false,
   "search.useIgnoreFiles": true
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,9 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.38.0] - 2025-06-24
+
 ### Added
 
 - `tilebox-datasets`: Added `delete_collection` method to `DatasetClient` to delete a collection.
+- `tilebox-workflows`: Added support for nested dataclasses and protobuf messages inside a `list`, `dict` or `tuple`
+  as task arguments.
 
 ## [0.37.1] - 2025-06-10
 
@@ -194,7 +198,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Released packages: `tilebox-datasets`, `tilebox-workflows`, `tilebox-storage`, `tilebox-grpc`
 
 
-[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.37.1...HEAD
+[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.38.0...HEAD
+[0.38.0]: https://github.com/tilebox/tilebox-python/compare/v0.37.1...v0.38.0
 [0.37.1]: https://github.com/tilebox/tilebox-python/compare/v0.37.0...v0.37.1
 [0.37.0]: https://github.com/tilebox/tilebox-python/compare/v0.36.1...v0.37.0
 [0.36.1]: https://github.com/tilebox/tilebox-python/compare/v0.36.0...v0.36.1
diff --git a/README.md b/README.md
@@ -31,7 +31,7 @@
 
 # Tilebox Python
 
-Python clients for [Tilebox](https://tilebox.com) - a framework for space data management and workflow orchestration.
+Python library for [Tilebox](https://tilebox.com), a lightweight space data management and orchestration software - on ground and in orbit.
 
 ## Install
 
diff --git a/tilebox-datasets/tests/data/test_time_interval.py b/tilebox-datasets/tests/data/test_time_interval.py
@@ -44,6 +44,12 @@ def test_time_interval_to_half_open_preserves_eq(interval: TimeInterval) -> None
     assert interval.to_half_open() == interval
 
 
+@given(time_intervals())
+def test_time_interval_to_half_open_preserves_hash(interval: TimeInterval) -> None:
+    """Assert that converting a time interval to a half-open interval preserves hash equality"""
+    assert hash(interval.to_half_open()) == hash(interval)
+
+
 @given(time_intervals())
 def test_time_interval_repr(interval: TimeInterval) -> None:
     """Assert the repr of a time interval is as expected"""
diff --git a/tilebox-datasets/tests/testdata/recordings/find_s2_datapoint.rpcs.bin b/tilebox-datasets/tests/testdata/recordings/find_s2_datapoint.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8072337bacb7142d3c501c7bc6be339d6bb84c3528ce524896d167bf9af36648
+oid sha256:5234a799245c656e37e16e249f6212b0c7a48021d9a2fbc7a672375d7354a57c
 size 10004
diff --git a/tilebox-datasets/tests/testdata/recordings/list_datasets.rpcs.bin b/tilebox-datasets/tests/testdata/recordings/list_datasets.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4d42759c7a576103c0f27b7ca099fa5bee57e82f1132b61773e20ec30a58b907
-size 51373
+oid sha256:b090da550257de86055244f03cedb98cb09a88a3551664ab9d8ebbd8641e3441
+size 57645
diff --git a/tilebox-datasets/tests/testdata/recordings/list_s2_collections.rpcs.bin b/tilebox-datasets/tests/testdata/recordings/list_s2_collections.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:051d73203c97c307b17339d6f40e878684c979286ccfd57c659211debd41954b
+oid sha256:0f928b7432ffedbd3896a89a7ff5e714eae7be5a4b01bc217bfc5fd6799abff8
 size 9154
diff --git a/tilebox-datasets/tests/testdata/recordings/query_sentinel2.rpcs.bin b/tilebox-datasets/tests/testdata/recordings/query_sentinel2.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8ed6af9697d66efc75fc996a6bcd37639716f458d8ddee7373e4bd29ae7e9994
+oid sha256:ef04121a4e99f85a25513d0933823b1c271a367b575dece8959be080c85e456d
 size 903956
diff --git a/tilebox-datasets/tests/testdata/recordings/query_sentinel2_paging.rpcs.bin b/tilebox-datasets/tests/testdata/recordings/query_sentinel2_paging.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fdb9d76547b77c03c477faa991411d5a3bd0ee1887e350be7400e59e777f84dc
+oid sha256:a5fee0adc442645fb623f35d133681f34003690f582d532f486ecd91f59cc67b
 size 873354
diff --git a/tilebox-datasets/tests/testdata/recordings/s2_collection_info.rpcs.bin b/tilebox-datasets/tests/testdata/recordings/s2_collection_info.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:051d73203c97c307b17339d6f40e878684c979286ccfd57c659211debd41954b
+oid sha256:0f928b7432ffedbd3896a89a7ff5e714eae7be5a4b01bc217bfc5fd6799abff8
 size 9154
diff --git a/tilebox-datasets/tests/testdata/recordings/s2_datapoint_not_found.rpcs.bin b/tilebox-datasets/tests/testdata/recordings/s2_datapoint_not_found.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d053f62c6c5e35d4db9b05772b8d876195d509d31166a3ed7697662f1af5b4f8
+oid sha256:7428cc2c23ba142cdc148053f5edf9aa5d7f630cf0139c69c13701c516f78037
 size 8560
diff --git a/tilebox-datasets/tilebox/datasets/data/time_interval.py b/tilebox-datasets/tilebox/datasets/data/time_interval.py
@@ -73,6 +73,13 @@ def __eq__(self, other: object) -> bool:
         a, b = self.to_half_open(), other.to_half_open()
         return (a.start, a.end) == (b.start, b.end)
 
+    def __hash__(self) -> int:
+        """Hash the time interval"""
+
+        # if two intervals are equal, they should have the same hash, so we convert to half-open intervals first
+        half_open = self.to_half_open()
+        return hash((half_open.start, half_open.end))
+
     def __repr__(self) -> str:
         return self.format()
 
diff --git a/tilebox-workflows/tests/runner/test_runner.py b/tilebox-workflows/tests/runner/test_runner.py
@@ -57,10 +57,10 @@ def test_runner_with_fibonacci_workflow() -> None:
         # we hardcode the trace parent for the job, which allows us to assert that every single outgoing request
         # matches exactly byte for byte
         get_trace_parent_mock.return_value = "00-f8d3b65869f638c5bfe173ffb3b3e5a0-ccf5709467cafc52-01"
-        job = client.jobs().submit("fibonacci", FibonacciTask(n), cluster="dev-cluster-ESugE7S4cwADVK")
+        job = client.jobs().submit("fibonacci", FibonacciTask(n))
 
     cache = InMemoryCache()
-    runner = client.runner("dev-cluster-ESugE7S4cwADVK", tasks=[FibonacciTask, SumResultTask], cache=cache)
+    runner = client.runner(tasks=[FibonacciTask, SumResultTask], cache=cache)
     runner.run_all()
 
     job_cache = cache.group(str(job.id))
@@ -85,10 +85,10 @@ def test_runner_with_flaky_task() -> None:
         # we hardcode the trace parent for the job, which allows us to assert that every single outgoing request
         # matches exactly byte for byte
         get_trace_parent_mock.return_value = "00-9680c9bfd602c4befe7b65a33a7b886d-3de78304f4cfbc40-01"
-        job = client.jobs().submit("flaky-task", FlakyTask(), cluster="dev-cluster-ESugE7S4cwADVK")
+        job = client.jobs().submit("flaky-task", FlakyTask())
 
     cache = InMemoryCache()
-    runner = client.runner("dev-cluster-ESugE7S4cwADVK", tasks=[FlakyTask], cache=cache)
+    runner = client.runner(tasks=[FlakyTask], cache=cache)
 
     runner.run_all()  # task will fail
     job = job_client.find(job)  # load current job state
diff --git a/tilebox-workflows/tests/runner/testdata/recordings/fibonacci_workflow.rpcs.bin b/tilebox-workflows/tests/runner/testdata/recordings/fibonacci_workflow.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:43960f4cd0782472294024a06cae56b1c77352be99f47b784a22d7cb7cdcc768
-size 9519
+oid sha256:8b2943e452addd5eedf0549a582b97b01ca746133a0f8171baa184e7c4dc0edf
+size 9244
diff --git a/tilebox-workflows/tests/runner/testdata/recordings/flaky_task.rpcs.bin b/tilebox-workflows/tests/runner/testdata/recordings/flaky_task.rpcs.bin
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f7bd111411d39ad34dfec99dd45fd64bf630395d32260ef08a677790e18aa11a
-size 5629
+oid sha256:375301ced6092e5ed6b1e40307fdaa14ddb34aa5637847750935e8d307a64050
+size 5489
diff --git a/tilebox-workflows/tests/test_task.py b/tilebox-workflows/tests/test_task.py
@@ -196,7 +196,7 @@ def test_serialize_deserialize_task_protobuf() -> None:
     assert deserialize_task(ExampleProtobufTaskWithSingleProtobufArg, serialize_task(task)) == task
 
 
-@dataclass
+@dataclass(frozen=True, eq=True)
 class NestedJson:
     nested_x: str
 
@@ -225,3 +225,126 @@ class ExampleTaskWithNestedProtobuf(Task):
 def test_serialize_deserialize_task_nested_protobuf() -> None:
     task = ExampleTaskWithNestedProtobuf("Hello", SampleArgs(some_string="World", some_int=123))
     assert deserialize_task(ExampleTaskWithNestedProtobuf, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedJsonInList(Task):
+    x: str
+    nested_list: list[NestedJson]
+
+
+def test_serialize_deserialize_task_nested_json_in_list() -> None:
+    task = ExampleTaskWithNestedJsonInList("Hello", [NestedJson("World"), NestedJson("!")])
+    assert deserialize_task(ExampleTaskWithNestedJsonInList, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedJsonInTuple(Task):
+    x: str
+    nested_tuple: tuple[NestedJson, NestedJson]
+
+
+def test_serialize_deserialize_task_nested_json_in_tuple() -> None:
+    task = ExampleTaskWithNestedJsonInTuple("Hello", (NestedJson("World"), NestedJson("!")))
+    assert deserialize_task(ExampleTaskWithNestedJsonInTuple, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedJsonInVariadicTuple(Task):
+    x: str
+    nested_tuple: tuple[NestedJson, ...]
+
+
+def test_serialize_deserialize_task_nested_json_in_variadic_tuple() -> None:
+    task = ExampleTaskWithNestedJsonInVariadicTuple("Hello", (NestedJson("World"), NestedJson("!"), NestedJson("!!!")))
+    assert deserialize_task(ExampleTaskWithNestedJsonInVariadicTuple, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedJsonListOnly(Task):
+    nested_list: list[NestedJson]
+
+
+def test_serialize_deserialize_task_nested_json_list_only() -> None:
+    task = ExampleTaskWithNestedJsonListOnly([NestedJson("World"), NestedJson("!")])
+    assert deserialize_task(ExampleTaskWithNestedJsonListOnly, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedJsonInDict(Task):
+    x: str
+    nested_list: dict[str, NestedJson]
+
+
+def test_serialize_deserialize_task_nested_json_in_dict() -> None:
+    task = ExampleTaskWithNestedJsonInDict("Hello", {"a": NestedJson("World"), "b": NestedJson("!")})
+    assert deserialize_task(ExampleTaskWithNestedJsonInDict, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedJsonInNestedDict(Task):
+    x: str
+    nested_list: dict[str, dict[str, list[NestedJson]]]
+
+
+def test_serialize_deserialize_task_nested_json_in_nested_dict() -> None:
+    task = ExampleTaskWithNestedJsonInNestedDict("Hello", {"a": {"b": [NestedJson("World"), NestedJson("!")]}})
+    assert deserialize_task(ExampleTaskWithNestedJsonInNestedDict, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedProtobufInList(Task):
+    x: str
+    nested_list: list[SampleArgs]
+
+
+def test_serialize_deserialize_task_nested_protobuf_in_list() -> None:
+    task = ExampleTaskWithNestedProtobufInList("Hello", [SampleArgs(some_string="World", some_int=123)])
+    assert deserialize_task(ExampleTaskWithNestedProtobufInList, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedProtobufInTuple(Task):
+    x: str
+    nested_tuple: tuple[SampleArgs, SampleArgs]
+
+
+def test_serialize_deserialize_task_nested_protobuf_in_tuple() -> None:
+    task = ExampleTaskWithNestedProtobufInTuple(
+        "Hello", (SampleArgs(some_string="World", some_int=123), SampleArgs(some_string="!", some_int=456))
+    )
+    assert deserialize_task(ExampleTaskWithNestedProtobufInTuple, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedProtobufInVariadicTuple(Task):
+    x: str
+    nested_tuple: tuple[SampleArgs, ...]
+
+
+def test_serialize_deserialize_task_nested_protobuf_in_variadic_tuple() -> None:
+    task = ExampleTaskWithNestedProtobufInVariadicTuple(
+        "Hello",
+        (
+            SampleArgs(some_string="World", some_int=123),
+            SampleArgs(some_string="!", some_int=456),
+            SampleArgs(some_string="!!!", some_int=789),
+        ),
+    )
+    assert deserialize_task(ExampleTaskWithNestedProtobufInVariadicTuple, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedProtobufInDict(Task):
+    x: str
+    nested_dict: dict[str, SampleArgs]
+
+
+def test_serialize_deserialize_task_nested_protobuf_in_dict() -> None:
+    task = ExampleTaskWithNestedProtobufInDict(
+        "Hello", {"a": SampleArgs(some_string="World", some_int=123), "b": SampleArgs(some_string="!", some_int=456)}
+    )
+    assert deserialize_task(ExampleTaskWithNestedProtobufInDict, serialize_task(task)) == task
+
+
+class ExampleTaskWithNestedProtobufInNestedDict(Task):
+    x: str
+    nested_dict: dict[str, dict[str, list[SampleArgs]]]
+
+
+def test_serialize_deserialize_task_nested_protobuf_in_nested_dict() -> None:
+    task = ExampleTaskWithNestedProtobufInNestedDict(
+        "Hello",
+        {"a": {"b": [SampleArgs(some_string="World", some_int=123), SampleArgs(some_string="!", some_int=456)]}},
+    )
+    assert deserialize_task(ExampleTaskWithNestedProtobufInNestedDict, serialize_task(task)) == task
diff --git a/tilebox-workflows/tilebox/workflows/task.py b/tilebox-workflows/tilebox/workflows/task.py
@@ -5,7 +5,7 @@
 from base64 import b64decode, b64encode
 from collections.abc import Sequence
 from dataclasses import dataclass, field, fields, is_dataclass
-from typing import Any, cast
+from typing import Any, cast, get_args, get_origin
 
 # from python 3.11 onwards this is available as typing.dataclass_transform:
 from typing_extensions import dataclass_transform
@@ -299,7 +299,7 @@ def serialize_task(task: Task) -> bytes:
         return b""  # empty task
     if len(task_fields) == 1:
         # if there is only one field, we can serialize it directly
-        field = _serialize_field(task, task_fields[0].name)
+        field = _serialize_value(getattr(task, task_fields[0].name), base64_encode_protobuf=False)
         if not isinstance(field, bytes):
             field = json.dumps(field).encode()
         return field
@@ -314,18 +314,24 @@ def _serialize_as_dict(task: Task) -> dict[str, Any]:
         if skip:
             continue
 
-        value = _serialize_field(task, dataclass_field.name)
-        if isinstance(value, bytes):
-            as_dict[dataclass_field.name] = b64encode(value).decode("ascii")
-        else:
-            as_dict[dataclass_field.name] = value
+        as_dict[dataclass_field.name] = _serialize_value(
+            getattr(task, dataclass_field.name), base64_encode_protobuf=True
+        )
 
     return as_dict
 
 
-def _serialize_field(task: Task, field_name: str) -> Any:
-    value = getattr(task, field_name)
+def _serialize_value(value: Any, base64_encode_protobuf: bool) -> Any:  # noqa: PLR0911
+    if isinstance(value, list):
+        return [_serialize_value(v, base64_encode_protobuf) for v in value]
+    if isinstance(value, tuple):
+        return tuple(_serialize_value(v, base64_encode_protobuf) for v in value)
+    if isinstance(value, dict):
+        # avoid serializing the dict keys, since nested dicts are not valid keys in dicts
+        return {k: _serialize_value(v, base64_encode_protobuf) for k, v in value.items()}
     if hasattr(value, "SerializeToString"):  # protobuf message
+        if base64_encode_protobuf:
+            return b64encode(value.SerializeToString()).decode("ascii")
         return value.SerializeToString()
     if is_dataclass(value):
         return _serialize_as_dict(value)  # type: ignore[arg-type]
@@ -348,9 +354,7 @@ def deserialize_task(task_cls: type, task_input: bytes) -> Task:
         if hasattr(field_type, "FromString"):  # protobuf message
             value = field_type.FromString(task_input)  # type: ignore[arg-type]
         else:
-            value = json.loads(task_input.decode())
-            if is_dataclass(field_type) and isinstance(value, dict):
-                value = _deserialize_dataclass(field_type, value)  # type: ignore[arg-type]
+            value = _deserialize_value(field_type, json.loads(task_input.decode()))  # type: ignore[arg-type]
 
         return task_cls(**{task_fields[0].name: value})
 
@@ -362,9 +366,35 @@ def _deserialize_dataclass(cls: type, params: dict[str, Any]) -> Task:
     for param in list(params):
         # recursively deserialize nested dataclasses
         field = cls.__dataclass_fields__[param]
-        if hasattr(field.type, "FromString"):
-            params[field.name] = field.type.FromString(b64decode(params[field.name]))
-        elif is_dataclass(field.type) and isinstance(params[field.name], dict):
-            params[field.name] = _deserialize_dataclass(field.type, params[field.name])  # type: ignore[arg-type]
+        params[field.name] = _deserialize_value(field.type, params[field.name])
 
     return cls(**params)
+
+
+def _deserialize_value(field_type: type, value: Any) -> Any:  # noqa: PLR0911
+    if hasattr(field_type, "FromString"):
+        return field_type.FromString(b64decode(value))
+    if is_dataclass(field_type) and isinstance(value, dict):
+        return _deserialize_dataclass(field_type, value)
+
+    # in case our field type is a list or dict, we need to recursively deserialize the values
+    origin_type = get_origin(field_type)
+    if not origin_type:
+        return value  # simple type, no further recursion needed
+
+    type_args = get_args(field_type)  # the wrapped type in a container, e.g. list[str] -> type_args is (str,)
+
+    if isinstance(value, list) and origin_type is list and len(type_args) == 1:
+        return [_deserialize_value(type_args[0], v) for v in value]
+    if isinstance(value, list) and origin_type is tuple:
+        # tuples are serialized as json list, so we get a list back
+        # which we want to convert back to a tuple
+        if len(type_args) == 2 and type_args[1] is Ellipsis:
+            type_args = (type_args[0],)  # variadic tuple, we only have one type argument to use for all values
+            return tuple(_deserialize_value(type_args[0], v) for v in value)
+        return tuple(_deserialize_value(type_args[min(i, len(type_args) - 1)], v) for i, v in enumerate(value))
+
+    if isinstance(value, dict) and origin_type is dict and len(type_args) == 2:
+        return {k: _deserialize_value(type_args[1], v) for k, v in value.items()}
+
+    return value
diff --git a/uv.lock b/uv.lock
