Optionally delete points layers after finishing alignment (#88)

Closes #9 
Closes #29 

---------

Co-authored-by: Juan Nunez-Iglesias <[email protected]>

src/affinder/_tests/test_affinder.py

@@ -165,3 +165,36 @@ def test_load_affine(tmp_path):
     widget(layer, affile)
 
     np.testing.assert_allclose(layer.affine, affine)
+
+
+@pytest.mark.parametrize('remove_pts', [True, False])
+def test_remove_points_layers(remove_pts, make_napari_viewer):
+    """Check whether remove_points_layer option actually removes the layers."""
+    ref_im = np.random.random((5, 5))
+    mov_im = np.random.random((5, 5))
+    ref_pts = np.array([[1, 1], [2, 2], [1, 4]], dtype=float)
+    mov_pts = np.array([[4, 1], [2, 2], [1, 4]], dtype=float)
+
+    viewer = make_napari_viewer()
+    ref_layer = viewer.add_image(ref_im)
+    mov_layer = viewer.add_image(mov_im)
+    qtwidget, widget = viewer.window.add_plugin_dock_widget(
+            'affinder', 'Start affinder'
+            )
+    widget(
+            viewer=viewer,
+            reference=ref_layer,
+            moving=mov_layer,
+            model=AffineTransformChoices.affine,
+            )
+    viewer.layers['ref_im_pts'].data = ref_pts
+    viewer.layers['mov_im_pts'].data = mov_pts
+
+    widget(delete_pts=remove_pts)  # close the widget
+
+    assert widget._call_button.text == 'Start'
+
+    assert remove_pts != any(
+            pt_layer in viewer.layers
+            for pt_layer in ['ref_im_pts', 'mov_im_pts']
+            )

src/affinder/affinder.py

@@ -1,11 +1,12 @@
+import functools
+import pathlib
 import warnings
+from enum import Enum
 from typing import Optional
 
-from enum import Enum
-import pathlib
-import toolz as tz
-from magicgui import magicgui, magic_factory
 import numpy as np
+import toolz as tz
+from magicgui import magic_factory
 from skimage.transform import (
         AffineTransform,
         EuclideanTransform,
@@ -78,14 +79,18 @@ def next_layer_callback(
             reset_view(viewer, reference_image_layer)
 
 
-# make a bindable function to shut things down
-@magicgui
 def close_affinder(layers, callback):
     for layer in layers:
         layer.events.data.disconnect(callback)
         layer.mode = 'pan_zoom'
 
 
+# make function to remove points layers after finishing
+def remove_pts_layers(viewer, layers):
+    for layer in layers:
+        viewer.layers.remove(layer)
+
+
 def _update_unique_choices(widget, choice_name):
     """Update the selected choice in a ComboBox widget to be unique.
 
@@ -126,6 +131,14 @@ def _on_affinder_main_init(widget):
         layout='vertical',
         output={'mode': 'w'},
         viewer={'visible': False, 'label': ' '},
+        delete_pts={
+                'label':
+                        'Delete points layers when done',
+                'tooltip': (
+                        'If ticked, the points layers used in alignment '
+                        'will be deleted when clicking "Finish".'
+                        ),
+                },
         )
 def start_affinder(
         viewer: 'napari.viewer.Viewer',
@@ -136,6 +149,7 @@ def start_affinder(
         moving_points: Optional['napari.layers.Points'] = None,
         model: AffineTransformChoices,
         output: Optional[pathlib.Path] = None,
+        delete_pts: bool = False,
         ):
     mode = start_affinder._call_button.text  # can be "Start" or "Finish"
 
@@ -181,13 +195,18 @@ def start_affinder(
         viewer.layers.selection.active = pts_layer0
         pts_layer0.mode = 'add'
 
-        close_affinder.layers.bind(points_layers)
-        close_affinder.callback.bind(callback)
-
+        start_affinder.close = functools.partial(
+                close_affinder, points_layers, callback
+                )
+        start_affinder.remove_points_layers = functools.partial(
+                remove_pts_layers, viewer, points_layers
+                )
         # change the button/mode for next run
         start_affinder._call_button.text = 'Finish'
     else:  # we are in Finish mode
-        close_affinder()
+        start_affinder.close()
+        if delete_pts:
+            start_affinder.remove_points_layers()
         start_affinder._call_button.text = 'Start'