Add handler to persist closing warnings

Author: f321x

electrum/gui/qt/main_window.py

@@ -88,7 +88,7 @@
                    filename_field, address_field, char_width_in_lineedit, webopen,
                    TRANSACTION_FILE_EXTENSION_FILTER_ANY, MONOSPACE_FONT,
                    getOpenFileName, getSaveFileName, ShowQRLineEdit, QtEventListener, qt_event_listener,
-                   event_listener, scan_qr_from_screenshot)
+                   event_listener, scan_qr_from_screenshot, QtClosingWarningHandler)
 from .wizard.wallet import WIF_HELP_TEXT
 from .history_list import HistoryList, HistoryModel
 from .update_checker import UpdateCheck, UpdateCheckThread
@@ -271,8 +271,7 @@ def add_optional_tab(tabs, tab, icon, description):
 
         # network callbacks
         self.register_callbacks()
-        # wallet closing warning callbacks
-        self.closing_warning_callbacks = []  # type: List[Callable[[], Optional[str]]]
+        self.closing_warning_handler = QtClosingWarningHandler(wallet)
         # banner may already be there
         if self.network and self.network.banner:
             self.console.showMessage(self.network.banner)
@@ -2690,7 +2689,7 @@ def _show_closing_warnings(self) -> bool:
         """Show any closing warnings and return True if the user chose to quit anyway."""
 
         warnings: Set[str] = set()
-        for cb in self.closing_warning_callbacks:
+        for cb in self.closing_warning_handler.get_closing_warning_callbacks():
             if warning := cb():
                 warnings.add(warning)
 
@@ -2710,23 +2709,6 @@ def _show_closing_warnings(self) -> bool:
             return True
         return False
 
-    def register_closing_warning_callback(self, callback: Callable[[], Optional[str]]) -> None:
-        """
-        Registers a callback that will be called when the wallet is closed. If the callback
-        returns a string it will be shown to the user as a warning to prevent them closing the wallet.
-        """
-        assert not asyncio.iscoroutinefunction(callback)
-        if not self.config.get("cmd") == "gui":
-            return
-        def warning_callback() -> Optional[str]:
-            try:
-                return callback()
-            except Exception:
-                self.logger.exception("Error in closing warning callback")
-                return None
-        self.logger.debug(f"registering wallet closing warning callback")
-        self.closing_warning_callbacks.append(warning_callback)
-
     def closeEvent(self, event):
         # note that closeEvent is NOT called if the user quits with Ctrl-C
         if not self._show_closing_warnings():
@@ -2918,13 +2900,12 @@ def on_swap_result(self, txid: Optional[str], *, is_reverse: bool):
             else:
                 msg += messages.MSG_FORWARD_SWAP_FUNDING_MEMPOOL + messages.MSG_FORWARD_SWAP_HTLC_EXPIRY
 
-            def wallet_closing_warning_callback() -> Optional[str]:
-                # gets called when the wallet GUI is closed
-                warning = messages.MSG_REVERSE_SWAP_FUNDING_MEMPOOL if is_reverse \
-                    else messages.MSG_FORWARD_SWAP_FUNDING_MEMPOOL
-                if self.wallet.adb.get_tx_height(txid).height < 1:
-                    return _("Ongoing submarine swap") + ":\n" + warning
-            self.register_closing_warning_callback(wallet_closing_warning_callback)
+            warning_msg = messages.MSG_REVERSE_SWAP_FUNDING_MEMPOOL if is_reverse \
+                        else messages.MSG_FORWARD_SWAP_FUNDING_MEMPOOL
+            self.closing_warning_handler.register_unconfirmed_tx_warning(
+                txid=txid,
+                warning_msg=_("Ongoing submarine swap") + ":\n" + warning_msg,
+            )
 
             self.show_message_signal.emit(msg)
         else:

electrum/gui/qt/util.py

@@ -7,6 +7,7 @@
 import os
 import webbrowser
 from functools import partial, lru_cache, wraps
+from asyncio import iscoroutinefunction
 from typing import (NamedTuple, Callable, Optional, TYPE_CHECKING, List, Any, Sequence, Tuple, Union)
 
 from PyQt6 import QtCore
@@ -20,6 +21,7 @@
                              QFrame, QAbstractButton)
 
 from electrum.i18n import _
+from electrum.logging import Logger
 from electrum.util import (FileImportFailed, FileExportFailed, resource_path, EventListener, event_listener,
                            get_logger, UserCancelled, UserFacingException)
 from electrum.invoices import (PR_UNPAID, PR_PAID, PR_EXPIRED, PR_INFLIGHT, PR_UNKNOWN, PR_FAILED, PR_ROUTING,
@@ -34,7 +36,7 @@
 
     from electrum.simple_config import SimpleConfig
     from electrum.simple_config import ConfigVarWithConfig
-
+    from electrum.wallet import Abstract_Wallet
 
 if platform.system() == 'Windows':
     MONOSPACE_FONT = 'Lucida Console'
@@ -748,6 +750,63 @@ def scan_qr_from_screenshot() -> QrCodeResult:
     assert len(scanned_qr) == 1, f"{len(scanned_qr)=}, expected 1"
     return scanned_qr[0]
 
+class QtClosingWarningHandler(Logger):
+    """
+    Provides handlers to register specific types of wallet closing warning callbacks,
+    and persists them.
+    """
+    def __init__(self, wallet: 'Abstract_Wallet'):
+        Logger.__init__(self)
+        self.wallet = wallet
+        persisted_callbacks = wallet.db.get_dict('wallet_closing_warning_callbacks')
+
+        # unconfirmed tx warnings: txid -> warning message, minimum confirmations
+        self.unconfirmed_tx_warnings: dict[str, Tuple[str, int]] \
+            = persisted_callbacks.setdefault('unconfirmed_tx_warnings', {})
+
+        self._closing_warning_callbacks = []  # type: List[Callable[[], Optional[str]]]
+        self._load_persisted_callbacks()
+
+    def _load_persisted_callbacks(self) -> None:
+        """Re-registers the callbacks loaded from the database."""
+        for txid, (warning_msg, minconf) in self.unconfirmed_tx_warnings.items():
+            self.register_unconfirmed_tx_warning(txid, warning_msg, minconf)
+
+    def _register_generic_warning_callback(
+        self,
+        callback: Callable[[], Optional[str]],
+        cleanup_cb: Optional[Callable[[], None]] = None
+    ) -> None:
+        """Wraps the callback in try/except and stores it in the callback list."""
+        assert not iscoroutinefunction(callback)
+        assert not iscoroutinefunction(cleanup_cb) if cleanup_cb else True
+
+        def warning_callback() -> Optional[str]:
+            warning: Optional[str] = None
+            try:
+                warning = callback()
+            except Exception:
+                self.logger.exception("Error in closing warning callback")
+            if warning is None and cleanup_cb:
+                self.logger.debug(f"removing closing warning callback {callback.__name__}")
+                cleanup_cb()
+            return warning
+
+        self.logger.debug(f"registering wallet closing warning callback: {callback.__name__}")
+        self._closing_warning_callbacks.append(warning_callback)
+
+    def register_unconfirmed_tx_warning(self, txid: str, warning_msg: str, minconf: int = 1) -> None:
+        """Registers a closing warning for a transaction that has to get confirmed."""
+        self.unconfirmed_tx_warnings[txid] = (warning_msg, minconf)
+        def unconfirmed_tx_warning_cb() -> Optional[str]:
+            if self.wallet.adb.get_tx_height(txid).height < minconf:
+                return warning_msg
+            return None
+        cleanup_cb = lambda: self.unconfirmed_tx_warnings.pop(txid, None)
+        self._register_generic_warning_callback(unconfirmed_tx_warning_cb, cleanup_cb)
+
+    def get_closing_warning_callbacks(self):
+        return self._closing_warning_callbacks
 
 class GenericInputHandler:
     def input_qr_from_camera(