Fix unit tests related to polygon/multiline tools

plotpy/tests/unit/test_annotation_tools.py

@@ -14,6 +14,7 @@
     AnnotatedEllipseTool,
     AnnotatedObliqueRectangleTool,
     AnnotatedPointTool,
+    AnnotatedPolygonTool,
     AnnotatedRectangleTool,
     AnnotatedSegmentTool,
     AverageCrossSectionTool,
@@ -44,6 +45,7 @@
     AnnotatedPointTool,
     AnnotatedRectangleTool,
     AnnotatedSegmentTool,
+    AnnotatedPolygonTool,
     AverageCrossSectionTool,
     CrossSectionTool,
     SnapshotTool,

plotpy/tests/unit/test_multiline_tools.py

@@ -14,78 +14,58 @@
 from plotpy.tools import MultiLineTool, PolygonTool
 
 
-def test_free_form_tool():
-    """Test the free form tool."""
-    corners = np.array(((0.1, 0.1), (0.1, 0.8), (0.8, 0.8), (0.8, 0.1)))
+def __generic_polyline_tool_test(
+    toolklass: type[MultiLineTool] | type[PolygonTool], points: np.ndarray
+) -> None:
+    """Generic polyline tool test."""
     with qt_app_context(exec_loop=False):
-        win, tool = create_window(PolygonTool)
+        win, tool = create_window(toolklass)
 
-        # drag_mouse(win, x_path, y_path)
-        for x, y in corners:
-            mouse_event_at_relative_plot_pos(win, (x, y), CLICK)
+        x0, y0 = None, None
+        for x, y in points:
+            if x0 is not None:
+                x_path = np.linspace(x0, x, 10)
+                y_path = np.linspace(y0, y, 10)
+                drag_mouse(win, x_path, y_path)
+            else:
+                mouse_event_at_relative_plot_pos(win, (x, y), CLICK)
+            x0, y0 = x, y
 
-        assert tool.shape is not None
-        assert tool.shape.get_points().shape == corners.shape
+        assert tool.handler.shape is not None
+        assert tool.handler.shape.get_points().shape == points.shape
 
-        # Delete last point
         keyboard_event(win, QC.Qt.Key.Key_Backspace)
 
-        points_count, _ = tool.shape.get_points().shape
+        points_count, _ = tool.handler.shape.get_points().shape
 
-        assert points_count == (len(corners) - 1)
+        assert points_count == (len(points) - 1)
 
         # add last point by dragging mouse
-        drag_mouse(win, corners[-2:, 0], corners[-2:, 1])
-
-        points_count, _ = tool.shape.get_points().shape
-        assert points_count == len(corners)
+        drag_mouse(win, points[-2:, 0], points[-2:, 1])
 
-        keyboard_event(win, QC.Qt.Key.Key_Enter)
+        points_count, _ = tool.handler.shape.get_points().shape
+        assert points_count == len(points)
 
-        assert tool.shape is None
+        keyboard_event(win, QC.Qt.Key.Key_Space)
 
         exec_dialog(win)
 
 
+def test_polygon_tool():
+    """Test the polygon tool."""
+    corners = np.array(((0.1, 0.1), (0.1, 0.8), (0.8, 0.8), (0.8, 0.1)))
+    __generic_polyline_tool_test(PolygonTool, corners)
+
+
 def test_multiline_tool():
     """Test the multi line tool."""
     n = 100
     t = np.linspace(0, np.pi * 10, n)
-
-    # Create x and y arrays
     x_arr = t * np.cos(t) / n + 0.5
     y_arr = t * np.sin(t) / n + 0.5
-
-    with qt_app_context(exec_loop=False):
-        win, tool = create_window(MultiLineTool)
-
-        # drag_mouse(win, x_path, y_path)
-        for x, y in zip(x_arr, y_arr):
-            mouse_event_at_relative_plot_pos(win, (x, y), CLICK)
-
-        assert tool.shape is not None
-        assert tool.shape.get_points().shape == np.array([x_arr, y_arr]).T.shape
-
-        # Delete last point
-        keyboard_event(win, QC.Qt.Key.Key_Backspace)
-
-        points_count, _ = tool.shape.get_points().shape
-
-        assert points_count == (n - 1)
-
-        # add last point by dragging mouse
-        drag_mouse(win, x_arr[-2:], y_arr[-2:])
-
-        points_count, _ = tool.shape.get_points().shape
-        assert points_count == n
-
-        keyboard_event(win, QC.Qt.Key.Key_Enter)
-
-        assert tool.shape is None
-
-        exec_dialog(win)
+    __generic_polyline_tool_test(MultiLineTool, np.array([x_arr, y_arr]).T)
 
 
 if __name__ == "__main__":
-    test_free_form_tool()
+    test_polygon_tool()
     test_multiline_tool()

plotpy/tools/shape.py

@@ -75,6 +75,7 @@ def __init__(
             tip=tip,
             switch_to_default_tool=switch_to_default_tool,
         )
+        self.handler: MultilineSelectionHandler | None = None
         self.switch_to_default_tool = switch_to_default_tool
         self.setup_shape_cb = setup_shape_cb
         self.handle_final_shape_cb = handle_final_shape_cb
@@ -127,12 +128,12 @@ def setup_filter(self, baseplot: BasePlot) -> StatefulEventFilter:
         """
         filter = baseplot.filter
         start_state = filter.new_state()
-        handler = MultilineSelectionHandler(
+        self.handler = MultilineSelectionHandler(
             filter, QC.Qt.LeftButton, start_state=start_state, closed=self.CLOSED
         )
-        handler.SIG_END_POLYLINE.connect(self.end_polyline)
+        self.handler.SIG_END_POLYLINE.connect(self.end_polyline)
         shape = self.get_shape()
-        handler.set_shape(shape, self.setup_shape)
+        self.handler.set_shape(shape, self.setup_shape)
         return setup_standard_tool_filter(filter, start_state)
 
     def end_polyline(self, filter: StatefulEventFilter, points: np.ndarray) -> None: