From 6d354b5f6ff312cb8069a2fa8bf6bf270f8391dc Mon Sep 17 00:00:00 2001
From: unexcellent <>
Date: Thu, 14 Nov 2024 05:20:29 +0100
Subject: [PATCH] feat: implement Poly3d.to_json()

---
 raillabel/format/poly3d.py                 | 14 +++++++++++++-
 tests/test_raillabel/format/test_poly3d.py |  5 +++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/raillabel/format/poly3d.py b/raillabel/format/poly3d.py
index 6c4777a..b49ff4e 100644
--- a/raillabel/format/poly3d.py
+++ b/raillabel/format/poly3d.py
@@ -8,7 +8,8 @@
 
 from raillabel.json_format import JSONPoly3d
 
-from ._attributes import _attributes_from_json
+from ._attributes import _attributes_from_json, _attributes_to_json
+from ._util import _flatten_list
 from .point3d import Point3d
 
 
@@ -45,6 +46,17 @@ def from_json(cls, json: JSONPoly3d, object_id: UUID) -> Poly3d:
             attributes=_attributes_from_json(json.attributes),
         )
 
+    def to_json(self, uid: UUID, object_type: str) -> JSONPoly3d:
+        """Export this object into the RailLabel JSON format."""
+        return JSONPoly3d(
+            name=self.name(object_type),
+            val=_flatten_list([point.to_json() for point in self.points]),
+            closed=self.closed,
+            coordinate_system=self.sensor_id,
+            uid=uid,
+            attributes=_attributes_to_json(self.attributes),
+        )
+
     def name(self, object_type: str) -> str:
         """Return the name of the annotation used for indexing in the object data pointers."""
         return f"{self.sensor_id}__poly3d__{object_type}"
diff --git a/tests/test_raillabel/format/test_poly3d.py b/tests/test_raillabel/format/test_poly3d.py
index 6119d38..66288f5 100644
--- a/tests/test_raillabel/format/test_poly3d.py
+++ b/tests/test_raillabel/format/test_poly3d.py
@@ -63,5 +63,10 @@ def test_name(poly3d):
     assert actual == "lidar__poly3d__person"
 
 
+def test_to_json(poly3d, poly3d_json, poly3d_id):
+    actual = poly3d.to_json(poly3d_id, object_type="person")
+    assert actual == poly3d_json
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])