ENH: Display a dialog to save trans and fids before quitting the coreg app

mne/gui/_coreg.py

@@ -174,6 +174,11 @@ def _get_default(var, val):
         self._to_cf_t = None
         self._omit_hsp_distance = 0.0
         self._fiducials_file = None
+        self._trans_modified = False
+        self._mri_fids_modified = False
+        self._mri_scale_modified = False
+        self._accept_close_event = True
+        self._auto_cleanup = True
         self._fid_colors = tuple(
             DEFAULTS['coreg'][f'{key}_color'] for key in
             ('lpa', 'nasion', 'rpa'))
@@ -225,7 +230,7 @@ def _get_default(var, val):
         # setup the window
         self._renderer = _get_renderer(
             size=self._defaults["size"], bgcolor=self._defaults["bgcolor"])
-        self._renderer._window_close_connect(self._clean)
+        self._renderer._window_close_connect(self._close_callback)
         self._renderer.set_interaction(interaction)
         self._renderer._status_bar_initialize()
 
@@ -309,6 +314,10 @@ def _get_default(var, val):
             self._renderer.plotter.add_callback(
                 self._redraw, self._refresh_rate_ms)
         self._renderer.plotter.show_axes()
+        # initialization does not count as modification by the user
+        self._trans_modified = False
+        self._mri_fids_modified = False
+        self._mri_scale_modified = False
         if block and self._renderer._kind != 'notebook':
             _qt_app_exec(self._renderer.figure.store["app"])
 
@@ -432,6 +441,7 @@ def _set_scale_mode(self, mode):
         self._scale_mode = mode
 
     def _set_fiducial(self, value, coord):
+        self._mri_fids_modified = True
         fid = self._current_fiducial
         fid_idx = _map_fid_name_to_idx(name=fid)
 
@@ -442,6 +452,10 @@ def _set_fiducial(self, value, coord):
         self._update_plot("mri_fids")
 
     def _set_parameter(self, value, mode_name, coord):
+        if mode_name == "scale":
+            self._mri_scale_modified = True
+        else:
+            self._trans_modified = True
         if self._params_locked:
             return
         if mode_name == "scale" and self._scale_mode == "uniform":
@@ -1222,6 +1236,7 @@ def _save_subject(self):
                 self._display_message(f"Computing {bem_name} solution..."
                                       " Done!")
         self._display_message(f"Saving {self._subject_to}... Done!")
+        self._mri_scale_modified = False
 
     def _save_mri_fiducials(self, fname):
         self._display_message(f"Saving {fname}...")
@@ -1231,11 +1246,13 @@ def _save_mri_fiducials(self, fname):
         )
         self._set_fiducials_file(fname)
         self._display_message(f"Saving {fname}... Done!")
+        self._mri_fids_modified = False
 
     def _save_trans(self, fname):
         write_trans(fname, self.coreg.trans, overwrite=True)
         self._display_message(
             f"{fname} transform file is saved.")
+        self._trans_modified = False
 
     def _load_trans(self, fname):
         mri_head_t = _ensure_trans(read_trans(fname, return_all=True),
@@ -1697,6 +1714,10 @@ def _configure_status_bar(self):
             'status_message', 'hide', value=None, input_value=False
         )
 
+    def _set_automatic_cleanup(self, state):
+        """Enable/Disable automatic cleanup (for testing purposes only)."""
+        self._auto_cleanup = state
+
     def _clean(self):
         self._renderer = None
         self._widgets.clear()
@@ -1709,3 +1730,68 @@ def _clean(self):
     def close(self):
         """Close interface and cleanup data structure."""
         self._renderer.close()
+
+    def _close_dialog_callback(self, button_name):
+        from ..viz.backends.renderer import MNE_3D_BACKEND_TESTING
+        self._accept_close_event = True
+        if button_name == "Save":
+            if self._trans_modified:
+                self._forward_widget_command(
+                    "save_trans", "set_value", None)
+                # cancel means _save_trans is not called
+                if self._trans_modified:
+                    self._accept_close_event = False
+            if self._mri_fids_modified:
+                self._forward_widget_command(
+                    "save_mri_fids", "set_value", None)
+            if self._mri_scale_modified:
+                if self._subject_to:
+                    self._save_subject()
+                else:
+                    dialog = self._renderer._dialog_warning(
+                        title="CoregistrationUI",
+                        text="The name of the output subject used to "
+                             "save the scaled anatomy is not set.",
+                        info_text="Please set a subject name",
+                        callback=lambda x: None,
+                        buttons=["Ok"],
+                        modal=not MNE_3D_BACKEND_TESTING,
+                    )
+                    dialog.show()
+                    self._accept_close_event = False
+        elif button_name == "Cancel":
+            self._accept_close_event = False
+        # This  should pass, but in case Windows has a different name for this
+        # than Linux and macOS, let's just skip it...
+        # else:
+        #     assert button_name in ("Discard", "Don't Save"), button_name
+
+    def _close_callback(self):
+        if self._trans_modified or self._mri_fids_modified or \
+                self._mri_scale_modified:
+            from ..viz.backends.renderer import MNE_3D_BACKEND_TESTING
+            # prepare the dialog's text
+            text = "The following is/are not saved:"
+            text += "<ul>"
+            if self._trans_modified:
+                text += "<li>Head&lt;&gt;MRI transform</li>"
+            if self._mri_fids_modified:
+                text += "<li>MRI fiducials</li>"
+            if self._mri_scale_modified:
+                text += "<li>scaled subject MRI</li>"
+            text += "</ul>"
+            self._widgets["close_dialog"] = self._renderer._dialog_warning(
+                title="CoregistrationUI",
+                text=text,
+                info_text="Do you want to save?",
+                callback=self._close_dialog_callback,
+                buttons=["Save", "Discard", "Cancel"],
+                # modal=True means that the dialog blocks the application
+                # when show() is called, until one of the buttons is clicked
+                modal=not MNE_3D_BACKEND_TESTING,
+            )
+            self._widgets["close_dialog"].show()
+
+        if self._accept_close_event and self._auto_cleanup:
+            self._clean()
+        return self._accept_close_event

mne/gui/tests/test_coreg_gui.py

@@ -146,6 +146,7 @@ def test_coreg_gui_pyvista(tmp_path, renderer_interactive_pyvistaqt):
     assert coreg._fiducials_file == fid_fname
 
     # fitting (with scaling)
+    assert not coreg._mri_scale_modified
     coreg._reset()
     coreg._reset_fitting_parameters()
     coreg._set_scale_mode("uniform")
@@ -161,19 +162,22 @@ def test_coreg_gui_pyvista(tmp_path, renderer_interactive_pyvistaqt):
                     atol=1e-3)
     coreg._set_scale_mode("None")
     coreg._set_icp_fid_match("matched")
+    assert coreg._mri_scale_modified
 
     # unlock fiducials
     assert coreg._lock_fids
     coreg._set_lock_fids(False)
     assert not coreg._lock_fids
 
     # picking
+    assert not coreg._mri_fids_modified
     vtk_picker = TstVTKPicker(coreg._surfaces['head'], 0, (0, 0))
     coreg._on_mouse_move(vtk_picker, None)
     coreg._on_button_press(vtk_picker, None)
     coreg._on_pick(vtk_picker, None)
     coreg._on_button_release(vtk_picker, None)
     coreg._on_pick(vtk_picker, None)  # also pick when locked
+    assert coreg._mri_fids_modified
 
     # lock fiducials
     coreg._set_lock_fids(True)
@@ -213,11 +217,18 @@ def test_coreg_gui_pyvista(tmp_path, renderer_interactive_pyvistaqt):
     assert coreg._head_resolution == \
         (config.get('MNE_COREG_HEAD_HIGH_RES', 'true') == 'true')
 
+    assert coreg._trans_modified
     tmp_trans = tmp_path / 'tmp-trans.fif'
     coreg._save_trans(tmp_trans)
+    assert not coreg._trans_modified
     assert op.isfile(tmp_trans)
 
+    # test _close_callback()
+    coreg._set_automatic_cleanup(False)
     coreg.close()
+    coreg._widgets['close_dialog'].trigger('Discard')  # do not save
+    coreg._clean()  # finally, cleanup internal structures
+
     # Coregistration instance should survive
     assert isinstance(coreg.coreg, Coregistration)
 

mne/viz/_brain/_brain.py

@@ -736,7 +736,7 @@ def _clean(self):
         for renderer in self._renderer._all_renderers:
             renderer.RemoveAllLights()
         # app_window cannot be set to None because it is used in __del__
-        for key in ('lighting', 'interactor', '_RenderWindow'):
+        for key in ('lighting', 'interactor'):
             setattr(self.plotter, key, None)
         # Qt LeaveEvent requires _Iren so we use _FakeIren instead of None
         # to resolve the ref to vtkGenericRenderWindowInteractor

mne/viz/backends/_qt.py

@@ -59,7 +59,12 @@ def _dialog_warning(self, title, text, info_text, callback, *,
         widget.setDefaultButton(default_button)
 
         def func(button):
-            callback(button.text())
+            # the text of the button may be prefixed by '&'
+            button_name = button.text().replace('&', '')
+            # handle MacOS Discard button
+            button_name = "Discard" \
+                if button_name == "Don't save" else button_name
+            callback(button_name)
 
         widget.buttonClicked.connect(func)
         return _QtDialogWidget(widget, modal)
@@ -539,13 +544,29 @@ def _window_initialize(self):
         self._window = self.figure.plotter.app_window
         self._window.setLocale(QLocale(QLocale.Language.English))
         self._window.signal_close.connect(self._window_clean)
+        self._window_close_callbacks = list()
+
+        # patch closeEvent
+        def closeEvent(event):
+            accept_close_event = True
+            for callback in self._window_close_callbacks:
+                ret = callback()
+                # check if one of the callbacks ignores the close event
+                if isinstance(ret, bool) and not ret:
+                    accept_close_event = False
+            if accept_close_event:
+                self._window.signal_close.emit()
+                event.accept()
+            else:
+                event.ignore()
+        self._window.closeEvent = closeEvent
 
     def _window_clean(self):
         self.figure._plotter = None
         self._interactor = None
 
     def _window_close_connect(self, func):
-        self._window.signal_close.connect(func)
+        self._window_close_callbacks.append(func)
 
     def _window_get_dpi(self):
         return self._window.windowHandle().screen().logicalDotsPerInch()