fix: make PointSet 2D and 3D classes iterable in Python (#96)

* fix: make PointSet 2D and 3D classes iterable in Python
open-space-collective · Sep 25, 2023 · 8de63a5 · 8de63a5
1 parent 79af18d
commit 8de63a5
Show file tree

Hide file tree

Showing 7 changed files with 204 additions and 157 deletions.
diff --git a/bindings/python/src/OpenSpaceToolkitMathematicsPy/Geometry/2D/Objects/PointSet.cpp b/bindings/python/src/OpenSpaceToolkitMathematicsPy/Geometry/2D/Objects/PointSet.cpp
@@ -40,8 +40,15 @@ inline void OpenSpaceToolkitMathematicsPy_Geometry_2D_Objects_PointSet(pybind11:
 
         .def_static("empty", &PointSet::Empty)
 
-        // .def("__iter__", boost::python::range(static_cast<PointSet::ConstIterator (PointSet::*)() const>
-        // (&PointSet::begin), static_cast<PointSet::ConstIterator (PointSet::*)() const> (&PointSet::end)))
+        .def("__len__", &PointSet::getSize)
+        .def(
+            "__iter__",
+            [](const PointSet& aPointSet)
+            {
+                return make_iterator(aPointSet.begin(), aPointSet.end());
+            },
+            keep_alive<0, 1>()
+        )  // Keep vector alive while iterator is used
 
         ;
 }
diff --git a/bindings/python/src/OpenSpaceToolkitMathematicsPy/Geometry/3D/Objects/LineString.cpp b/bindings/python/src/OpenSpaceToolkitMathematicsPy/Geometry/3D/Objects/LineString.cpp
@@ -42,7 +42,8 @@ inline void OpenSpaceToolkitMathematicsPy_Geometry_3D_Objects_LineString(pybind1
             {
                 return aLineString.accessPointAt(anIndex);
             },
-            return_value_policy::reference_internal
+            return_value_policy::reference_internal,
+            arg("index")
         )
         .def(
             "__iter__",

diff --git a/bindings/python/src/OpenSpaceToolkitMathematicsPy/Geometry/3D/Objects/PointSet.cpp b/bindings/python/src/OpenSpaceToolkitMathematicsPy/Geometry/3D/Objects/PointSet.cpp
@@ -33,8 +33,15 @@ inline void OpenSpaceToolkitMathematicsPy_Geometry_3D_Objects_PointSet(pybind11:
 
         .def_static("empty", &PointSet::Empty)
 
-        // .def_static("__iter__", boost::python::range(static_cast<PointSet::ConstIterator (PointSet::*)() const>
-        // (&PointSet::begin), static_cast<PointSet::ConstIterator (PointSet::*)() const> (&PointSet::end)))
+        .def("__len__", &PointSet::getSize)
+        .def(
+            "__iter__",
+            [](const PointSet& aPointSet)
+            {
+                return make_iterator(aPointSet.begin(), aPointSet.end());
+            },
+            keep_alive<0, 1>()
+        )  // Keep vector alive while iterator is used
 
         ;
 }
diff --git a/bindings/python/test/geometry/d2/objects/test_linestring.py b/bindings/python/test/geometry/d2/objects/test_linestring.py
@@ -1,14 +1,15 @@
 # Apache License 2.0
 
-import pytest
+from collections.abc import Iterator, Iterable
 
 import numpy as np
-from collections.abc import Iterator, Iterable
 
-import ostk.mathematics as mathematics
+import pytest
 
 from ostk.core.types import String
 
+import ostk.mathematics as mathematics
+
 
 Object = mathematics.geometry.d2.Object
 Point = mathematics.geometry.d2.objects.Point
@@ -124,20 +125,16 @@ def test_get_point_count_success(self):
         assert len(linestring) == linestring.get_point_count()
         assert len(LineString.empty()) == 0
 
-    def test_iter_success(self):
+    def test_len_success(self):
         point_1: Point = Point(-1.0, 1.0)
         point_2: Point = Point(1.0, 1.0)
 
         linestring: LineString = LineString([point_1, point_2])
 
-        for point in linestring:
-            assert isinstance(point, Point)
-
-        assert iter(linestring) is not None
-        assert isinstance(iter(linestring), Iterator)
-        assert isinstance(iter(linestring), Iterable)
+        assert len(linestring) == 2
+        assert len(LineString.empty()) == 0
 
-    def test_getitem(self):
+    def test_getitem_success(self):
         point_1: Point = Point(-1.0, 1.0)
         point_2: Point = Point(1.0, 1.0)
 
@@ -159,6 +156,19 @@ def test_getitem(self):
         assert isinstance(linestring_list[1], Point)
         assert len(linestring_list) == 2
 
-    # def test_to_string_success (self):
+    def test_iter_success(self):
+        point_1: Point = Point(-1.0, 1.0)
+        point_2: Point = Point(1.0, 1.0)
+
+        linestring: LineString = LineString([point_1, point_2])
+
+        for point in linestring:
+            assert isinstance(point, Point)
+
+        assert iter(linestring) is not None
+        assert isinstance(iter(linestring), Iterator)
+        assert isinstance(iter(linestring), Iterable)
+
+    # def test_to_string_success(self):
 
-    # def test_apply_transformation_success (self):
+    # def test_apply_transformation_success(self):
diff --git a/bindings/python/test/geometry/d2/objects/test_point_set.py b/bindings/python/test/geometry/d2/objects/test_point_set.py
@@ -1,5 +1,7 @@
 # Apache License 2.0
 
+from collections.abc import Iterator, Iterable
+
 import pytest
 
 import ostk.mathematics as mathematics
@@ -79,10 +81,22 @@ def test_get_size_success(self):
     def test_distance_to_success_point(self, point_set: PointSet):
         assert point_set.distance_to(Point(1.0, 2.0)) == 0.0
 
-    # def test_is_near_success (self):
+    def test_len_success(self, point_set: PointSet):
+        assert len(point_set) == 2
+        assert len(PointSet.empty()) == 0
+
+    def test_iter_success(self, point_set: PointSet):
+        for point in point_set:
+            assert isinstance(point, Point)
+
+        assert iter(point_set) is not None
+        assert isinstance(iter(point_set), Iterator)
+        assert isinstance(iter(point_set), Iterable)
+
+    # def test_is_near_success(self):
 
-    # def test_get_point_closest_to_success (self):
+    # def test_get_point_closest_to_success(self):
 
-    # def test_to_string_success (self):
+    # def test_to_string_success(self):
 
-    # def test_apply_transformation_success (self):
+    # def test_apply_transformation_success(self):