feat: implement SceneBuilder.add_object()

raillabel/scene_builder/scene_builder.py

@@ -3,18 +3,68 @@
 
 from __future__ import annotations
 
+from copy import deepcopy
 from dataclasses import dataclass
+from uuid import UUID
 
-from raillabel.format import Metadata, Scene
+from raillabel.format import Metadata, Object, Scene
 
 
 @dataclass
 class SceneBuilder:
     """Use this class for easily creating scenes for tests."""
 
-    scene: Scene
+    result: Scene
 
     @classmethod
     def empty(cls) -> SceneBuilder:
         """Construct the SceneBuilder with an empty scene."""
         return SceneBuilder(Scene(metadata=Metadata(schema_version="1.0.0")))
+
+    def add_object(
+        self,
+        object_id: str | UUID | None = None,
+        object_type: str | None = None,
+        object_name: str | None = None,
+    ) -> SceneBuilder:
+        """Add an object to a scene."""
+        scene = deepcopy(self.result)
+
+        object_type, object_name = _resolve_empty_object_name_or_type(object_type, object_name)
+        object_id = _resolve_empty_object_uid(scene, object_id)
+
+        scene.objects[object_id] = Object(object_name, object_type)
+        return SceneBuilder(scene)
+
+
+def _resolve_empty_object_name_or_type(
+    object_type: str | None, object_name: str | None
+) -> tuple[str, str]:
+    if object_name is None and object_type is None:
+        object_type = "person"
+        object_name = object_type + "_0000"
+        return object_type, object_name
+
+    if object_name is None and object_type is not None:
+        object_name = object_type + "_0000"
+        return object_type, object_name
+
+    if object_name is not None and object_type is None:
+        object_type = object_name.split("_")[0]
+        return object_type, object_name
+
+    return object_type, object_name  # type: ignore
+
+
+def _resolve_empty_object_uid(scene: Scene, object_id: str | UUID | None) -> UUID:
+    if object_id is None:
+        uid_index = 0
+        while _generate_deterministic_uuid(uid_index, "5c59aad4") in scene.objects:
+            uid_index += 1
+        object_id = _generate_deterministic_uuid(uid_index, "5c59aad4")
+
+    return UUID(str(object_id))
+
+
+def _generate_deterministic_uuid(index: int, prefix: str) -> UUID:
+    return UUID(f"{prefix.zfill(0)}-0000-4000-0000-{str(index).zfill(12)}")

tests/test_raillabel/scene_builder/test_scene_builder.py

@@ -2,17 +2,93 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
+from uuid import UUID
 
 import pytest
 
 import raillabel
-from raillabel.format import Scene, Metadata
+from raillabel.format import Scene, Metadata, Object
 from raillabel.scene_builder.scene_builder import SceneBuilder
 
 
 def test_empty():
     actual = SceneBuilder.empty()
-    assert actual.scene == Scene(Metadata(schema_version="1.0.0"))
+    assert actual.result == Scene(Metadata(schema_version="1.0.0"))
+
+
+def test_add_object__all_options():
+    actual = SceneBuilder.empty().add_object(
+        object_id="5c59aad4-9fcd-4903-a9fa-72b1b76c23a5",
+        object_type="train",
+        object_name="train_0001",
+    )
+    assert actual.result == Scene(
+        metadata=Metadata(schema_version="1.0.0"),
+        objects={
+            UUID("5c59aad4-9fcd-4903-a9fa-72b1b76c23a5"): Object(name="train_0001", type="train")
+        },
+    )
+
+
+def test_add_object__no_object_name():
+    actual = SceneBuilder.empty().add_object(
+        object_id="5c59aad4-9fcd-4903-a9fa-72b1b76c23a5",
+        object_type="train",
+    )
+    assert actual.result == Scene(
+        metadata=Metadata(schema_version="1.0.0"),
+        objects={
+            UUID("5c59aad4-9fcd-4903-a9fa-72b1b76c23a5"): Object(name="train_0000", type="train")
+        },
+    )
+
+
+def test_add_object__no_object_type():
+    actual = SceneBuilder.empty().add_object(
+        object_id="5c59aad4-9fcd-4903-a9fa-72b1b76c23a5",
+        object_name="train_0000",
+    )
+    assert actual.result == Scene(
+        metadata=Metadata(schema_version="1.0.0"),
+        objects={
+            UUID("5c59aad4-9fcd-4903-a9fa-72b1b76c23a5"): Object(name="train_0000", type="train")
+        },
+    )
+
+
+def test_add_object__no_object_type_or_name():
+    actual = SceneBuilder.empty().add_object(
+        object_id="5c59aad4-9fcd-4903-a9fa-72b1b76c23a5",
+    )
+    assert actual.result == Scene(
+        metadata=Metadata(schema_version="1.0.0"),
+        objects={
+            UUID("5c59aad4-9fcd-4903-a9fa-72b1b76c23a5"): Object(name="person_0000", type="person")
+        },
+    )
+
+
+def test_add_object__no_object_id():
+    actual = SceneBuilder.empty().add_object(object_type="train", object_name="train_0001")
+    assert actual.result == Scene(
+        metadata=Metadata(schema_version="1.0.0"),
+        objects={
+            UUID("5c59aad4-0000-4000-0000-000000000000"): Object(name="train_0001", type="train")
+        },
+    )
+
+
+def test_add_object__object_id_iteration():
+    actual = SceneBuilder.empty().add_object().add_object().add_object().add_object()
+    assert actual.result == Scene(
+        metadata=Metadata(schema_version="1.0.0"),
+        objects={
+            UUID("5c59aad4-0000-4000-0000-000000000000"): Object(name="person_0000", type="person"),
+            UUID("5c59aad4-0000-4000-0000-000000000001"): Object(name="person_0000", type="person"),
+            UUID("5c59aad4-0000-4000-0000-000000000002"): Object(name="person_0000", type="person"),
+            UUID("5c59aad4-0000-4000-0000-000000000003"): Object(name="person_0000", type="person"),
+        },
+    )
 
 
 if __name__ == "__main__":