Surface.add_rectangle: accept 3 or 4 vertex indices (previously only 3)

examples/precentral_gyrus_border.ipynb

@@ -88,12 +88,11 @@
     "    corner_indices = [[surface.add_vertex(Vertex(*(coords + forward_dir))) for coords in cross_section_corners]\n",
     "                      for forward_dir in [-forward_vector, +forward_vector]]\n",
     "    for surface_corner_indices in corner_indices:\n",
-    "        surface.triangles.append(Triangle(surface_corner_indices[:3]))\n",
-    "        surface.triangles.append(Triangle(surface_corner_indices[2:] + surface_corner_indices[:1]))\n",
+    "        surface.add_rectangle(surface_corner_indices)\n",
     "    for i in range(4):\n",
     "        j = (i + 1) % 4\n",
-    "        surface.triangles.append(Triangle((corner_indices[0][i], corner_indices[0][j], corner_indices[1][j])))\n",
-    "        surface.triangles.append(Triangle((corner_indices[0][i], corner_indices[1][j], corner_indices[1][i])))\n",
+    "        surface.add_rectangle((corner_indices[0][i], corner_indices[0][j],\n",
+    "                               corner_indices[1][j], corner_indices[1][i]))\n",
     "\n",
     "surface.write_triangular('precentral-border.lh.pial')"
    ]

freesurfer_surface/__init__.py

@@ -392,12 +392,13 @@ class Surface:
 
     def add_rectangle(self, vertex_indices: typing.Iterable[int]) -> typing.Iterable[int]:
         vertex_indices = list(vertex_indices)
-        assert len(vertex_indices) == 3
-        vertex_coords = [numpy.array(self.vertices[vertex_index])
-                         for vertex_index in vertex_indices]
-        vertex_coords.append(vertex_coords[0]
-                             + vertex_coords[2] - vertex_coords[1])
-        vertex_indices.append(self.add_vertex(Vertex(*vertex_coords[3])))
+        if len(vertex_indices) == 3:
+            vertex_indices.append(self.add_vertex(
+                self.vertices[vertex_indices[0]]
+                + self.vertices[vertex_indices[2]]
+                - self.vertices[vertex_indices[1]]
+            ))
+        assert len(vertex_indices) == 4
         self.triangles.append(Triangle(vertex_indices[:3]))
         self.triangles.append(Triangle(vertex_indices[2:]
                                        + vertex_indices[:1]))

tests/test_surface.py

@@ -153,6 +153,26 @@ def test_add_vertex():
     assert surface.vertices[1].right == pytest.approx(-3.0)
 
 
+@pytest.mark.parametrize('vertices_coords', [
+    ((0, 0, 0), (2, 4, 0), (2, 4, 3)),
+    ((0, 0, 0), (2, 4, 0), (2, 4, 3), (0, 0, 3)),
+    ((1, 1, 0), (3, 5, 0), (3, 5, 3), (1, 1, 3)),
+    ((1, 1, 7), (3, 5, 7), (3, 5, 3), (1, 1, 3)),
+    ((1, 1, 1), (3, 5, 7), (3, 5, 9), (1, 1, 3)),
+    ((3, 5, 7), (1, 1, 1), (1, 1, 3)),
+    ((3, 5, 7), (1, 1, 1), (1, 1, 3), (3, 5, 9)),
+])
+def test_add_rectangle(vertices_coords):
+    surface = Surface()
+    for vertex_coords in vertices_coords:
+        surface.add_vertex(Vertex(*(float(c) for c in vertex_coords)))
+    surface.add_rectangle(range(len(vertices_coords)))
+    assert len(surface.vertices) == 4
+    assert len(surface.triangles) == 2
+    assert surface.triangles[0].vertex_indices == (0, 1, 2)
+    assert surface.triangles[1].vertex_indices == (2, 3, 0)
+
+
 @pytest.mark.parametrize(('vertices_coords', 'expected_extra_vertex_coords'), [
     (((0, 0, 0), (2, 4, 0), (2, 4, 3)), (0, 0, 3)),
     (((1, 1, 0), (3, 5, 0), (3, 5, 3)), (1, 1, 3)),
@@ -160,16 +180,13 @@ def test_add_vertex():
     (((1, 1, 1), (3, 5, 7), (3, 5, 9)), (1, 1, 3)),
     (((3, 5, 7), (1, 1, 1), (1, 1, 3)), (3, 5, 9)),
 ])
-def test_add_rectangle(vertices_coords, expected_extra_vertex_coords):
+def test_add_rectangle_3(vertices_coords, expected_extra_vertex_coords):
     surface = Surface()
     for vertex_coords in vertices_coords:
         surface.add_vertex(Vertex(*(float(c) for c in vertex_coords)))
     surface.add_rectangle(range(3))
     assert tuple(surface.vertices[3]) \
         == pytest.approx(expected_extra_vertex_coords)
-    assert len(surface.triangles) == 2
-    assert surface.triangles[0].vertex_indices == (0, 1, 2)
-    assert surface.triangles[1].vertex_indices == (2, 3, 0)
 
 
 def test__get_vertex_label_index():