Add session switching

Node Mailer now lets you transfer user discovery to a different Nuke session instead of just displaying an empty user list!

Author: mervinvb

node_mailer/controller.py

@@ -19,7 +19,11 @@
 from .user_interface.about import AboutWindow
 from .user_interface.history import HistoryWindow
 from .user_interface.mailing import MailingWindow
-from .user_interface.popups import ReceivedMailPopup, display_error_popup
+from .user_interface.popups import (
+    ReceivedMailPopup,
+    display_error_popup,
+    display_yes_no_popup,
+)
 from .user_interface.settings import SettingsWindow
 
 
@@ -59,12 +63,24 @@ def connect_signals(self) -> None:
         self.history_window.import_mail.connect(nuke_interfacing.import_mail)
         self.mailing_window.send_mail.connect(self.send_mail)
         self.direct_messaging_model.message_received.connect(self.mail_received)
+        self.direct_messaging_model.shutdown_received.connect(self.shutdown_received)
 
     def open_mailing_window(self) -> None:
         """Opens the mailing window."""
-        self.mailing_window.show()
         play_click_sound()
 
+        if self.discovery_model.running:
+            self.mailing_window.show()
+            return
+
+        should_start = display_yes_no_popup(
+            "Node Mailer is not currently running in this Nuke session. Do you want to start it? This will disable Node Mailer in other running Nuke sessions."
+        )
+        if should_start:
+            self.direct_messaging_model.send_shutdown_message()
+            self.reinitialize_systems()
+            self.mailing_window.show()
+
     def open_history_window(self) -> None:
         """Opens the history window."""
         self.history_storage_model.retrieve_all_mail_from_database()
@@ -115,6 +131,15 @@ def send_mail(self, client: NodeMailerClient, message: str) -> None:
         self.mailing_window.message_text_edit.clear()
         self.mailing_window.close()
 
+    def reinitialize_systems(self) -> None:
+        """Kills other Node Mailer process if running and, after a slight delay,
+        reinitializes all systems/models required for Node Mailer to function.
+        Delay is there to give the other Node Mailer instance some time to shut down."""
+        self.direct_messaging_model.send_shutdown_message()
+        QtCore.QTimer.singleShot(1000, self.discovery_model.initialize_socket)
+        QtCore.QTimer.singleShot(1000, self.discovery_model.start_background_processes)
+        QtCore.QTimer.singleShot(1000, self.direct_messaging_model.start_listening)
+
     def mail_received(self, mail: NodeMailerMail) -> None:
         """Handles a received mail using the received mail popup.
 
@@ -132,3 +157,10 @@ def mail_received(self, mail: NodeMailerMail) -> None:
 
         if received_mail_popup.picked_option == ReceivedMailPopupOption.IMPORT:
             nuke_interfacing.import_mail(mail)
+
+    def shutdown_received(self) -> None:
+        """Shuts down Node Mailer systems when we receive a shutdown message so another
+        Node Mailer instance can bind to the ports again."""
+        self.discovery_model.uninitialize_socket()
+        self.direct_messaging_model.stop_listening()
+        self.mailing_window.close()

node_mailer/data_models.py

@@ -27,9 +27,9 @@ class NodeMailerMail:
     node_string: str
     timestamp: int
 
-    def as_json(self) -> str:
-        """Returns the message as a JSON string."""
-        return json.dumps(asdict(self))
+    def as_dict(self) -> dict:
+        """Returns the message as a dictionary."""
+        return asdict(self)
 
 
 @dataclass

node_mailer/models/discovery.py

@@ -21,6 +21,7 @@ class ClientDiscovery(QtCore.QAbstractListModel):
     def __init__(self) -> None:
         """Initializes the model class."""
         super().__init__()
+        self.running = False
         self.mailing_clients: List[NodeMailerClient] = []
 
         self.local_addresses = self.get_local_ip_addresses()
@@ -52,6 +53,14 @@ def initialize_socket(self) -> None:
             constants.Ports.BROADCAST.value,
         )
         self.udp_socket.readyRead.connect(self.on_datagram_received)
+        self.running = self.udp_socket.state() == QtNetwork.QAbstractSocket.BoundState
+
+    def uninitialize_socket(self) -> None:
+        """Uninitializes the socket and disconnects the signals."""
+        self.udp_socket.readyRead.disconnect(self.on_datagram_received)
+        self.udp_socket.close()
+        self.running = False
+        self.mailing_clients = []
 
     def store_icons(self) -> None:
         """Stores the icons this model should return."""

node_mailer/models/messaging.py

@@ -20,6 +20,7 @@ class DirectMessaging(QtCore.QObject):
     Uses TCP to receive/send data."""
 
     message_received = QtCore.Signal(NodeMailerMail)
+    shutdown_received = QtCore.Signal()
 
     def __init__(self) -> None:
         """Initializes the messaging handler."""
@@ -36,6 +37,10 @@ def start_listening(self) -> None:
             port=constants.Ports.MESSAGING.value,
         )
 
+    def stop_listening(self) -> None:
+        """Stops listening for incoming messages."""
+        self.tcp_server.close()
+
     def on_new_connection(self) -> None:
         """Connects the readyRead signal to the on_message_received slot for new connections."""
         new_client = self.tcp_server.nextPendingConnection()
@@ -44,7 +49,7 @@ def on_new_connection(self) -> None:
             lambda: self.on_message_received(receiving_connection)
         )
         new_client.disconnected.connect(
-            lambda: self.on_connection_closed(receiving_connection)
+            lambda: self.process_received_message(receiving_connection)
         )
         self.open_connections.append(receiving_connection)
 
@@ -56,27 +61,24 @@ def on_message_received(self, connection: ReceivingConnection) -> None:
         """
         connection.message += connection.socket.readAll().data().decode("utf-8")
 
-    def on_connection_closed(self, connection: ReceivingConnection) -> None:
-        """Closed connection means we have received all the data. We can now process the message.
+    def process_received_message(self, connection: ReceivingConnection) -> None:
+        """Processes the complete received message we have stored when the sending connection closes.
 
         Args:
             connection: The connection that was closed.
         """
-        mail = self.get_mail_from_message_string(connection.message)
-        self.message_received.emit(mail)
-        self.open_connections.remove(connection)
+        try:
+            parsed_message = json.loads(connection.message)
+        except json.JSONDecodeError:
+            return
 
-    def get_mail_from_message_string(self, message_string: str) -> NodeMailerMail:
-        """Returns a NodeMailerMessage object from the network-sent string.
-
-        Args:
-            message_string: The JSON message in string form.
+        if parsed_message["type"] == "shutdown":
+            self.shutdown_received.emit()
+            return
 
-        Returns:
-            The NodeMailerMessage object.
-        """
-        parsed_message = json.loads(message_string)
-        return NodeMailerMail(**parsed_message)
+        mail = NodeMailerMail(**parsed_message["mail"])
+        self.message_received.emit(mail)
+        self.open_connections.remove(connection)
 
     def send_mail_to_client(
         self, mail: NodeMailerMail, client: NodeMailerClient
@@ -90,10 +92,29 @@ def send_mail_to_client(
         tcp_socket = QtNetwork.QTcpSocket()
         tcp_socket.connectToHost(client.ip_address, constants.Ports.MESSAGING.value)
 
-        if not tcp_socket.waitForConnected(2000):
+        if not tcp_socket.waitForConnected(500):
             msg = f"Could not connect to client {client.name}. Is it still running?"
             raise ConnectionError(msg)
 
-        tcp_socket.write(mail.as_json().encode("utf-8"))
+        dict_mail = {"type": "mail"}
+        dict_mail["mail"] = mail.as_dict()
+
+        tcp_socket.write(json.dumps(dict_mail).encode("utf-8"))
         tcp_socket.waitForBytesWritten()
         tcp_socket.disconnectFromHost()
+        tcp_socket.close()
+
+    def send_shutdown_message(self) -> None:
+        """Sends a shutdown message to the local running Node Mailer instance."""
+        tcp_socket = QtNetwork.QTcpSocket()
+        tcp_socket.connectToHost("localhost", constants.Ports.MESSAGING.value)
+
+        if not tcp_socket.waitForConnected(2000):
+            return
+
+        shutdown_message = {"type": "shutdown"}
+        tcp_socket.write(json.dumps(shutdown_message).encode("utf-8"))
+        tcp_socket.waitForBytesWritten()
+        tcp_socket.disconnectFromHost()
+        tcp_socket.close()
+

node_mailer/resources/question.png

@@ -125,6 +125,136 @@ def display_error_popup(error_text: str) -> None:
     error_popup.exec_()
 
 
+class YesNoPopup(NodeMailerWindow):
+    """A popup window that displays a yes/no question."""
+
+    closed = QtCore.Signal()
+
+    def __init__(self, question_text: str) -> None:
+        """Initializes the yes/no popup.
+
+        Args:
+            question_text: The question to display.
+        """
+        super().__init__(
+            self.get_user_interface(question_text),
+            "Node Mailer: Question",
+            resizable=False,
+        )
+        self.picked_option = False
+
+    def get_user_interface(self, question_text: str) -> QtWidgets.QWidget:
+        """Returns the user interface for the yes/no popup.
+
+        Args:
+            question_text: The question to display.
+
+        Returns:
+            The widget with the yes/no popup user interface.
+        """
+        widget = QtWidgets.QWidget()
+        layout = QtWidgets.QVBoxLayout()
+        layout.setContentsMargins(0, 0, 0, 0)
+        widget.setLayout(layout)
+
+        layout.addWidget(self.get_question_widget(question_text))
+        layout.addWidget(self.get_bottom_buttons())
+
+        return widget
+
+    def get_question_widget(self, question_text: str) -> QtWidgets.QWidget:
+        """Returns the question widget for the yes/no popup.
+
+        Args:
+            question_text: The question to display.
+
+        Returns:
+            The widget with the question.
+        """
+        widget = QtWidgets.QWidget()
+        layout = QtWidgets.QHBoxLayout()
+        layout.setContentsMargins(20, 0, 20, 0)
+        layout.setAlignment(QtCore.Qt.AlignCenter)
+        layout.setSpacing(20)
+        widget.setLayout(layout)
+
+        question_icon = QtWidgets.QLabel()
+        pixmap = QtGui.QPixmap(
+            str(Path(__file__).parent.parent / "resources" / "question.png")
+        )
+        scaled_pixmap = pixmap.scaled(48, 48, QtCore.Qt.KeepAspectRatio)
+        question_icon.setPixmap(scaled_pixmap)
+        layout.addWidget(question_icon)
+
+        question_text = QtWidgets.QLabel(question_text)
+        question_text.setWordWrap(True)
+        question_text.setMinimumWidth(250)
+        layout.addWidget(question_text)
+
+        return widget
+
+    def get_bottom_buttons(self) -> QtWidgets.QWidget:
+        """Returns the bottom buttons for the yes/no popup.
+
+        Returns:
+            The widget with the bottom buttons.
+        """
+        widget = QtWidgets.QWidget()
+        layout = QtWidgets.QHBoxLayout()
+        layout.setContentsMargins(0, 0, 0, 0)
+        layout.setAlignment(QtCore.Qt.AlignRight)
+        widget.setLayout(layout)
+
+        no_button = NodeMailerButton("No")
+        no_button.clicked.connect(self.on_no_clicked)
+        layout.addWidget(no_button)
+
+        yes_button = NodeMailerButton("Yes")
+        yes_button.clicked.connect(self.on_yes_clicked)
+        layout.addWidget(yes_button)
+
+        return widget
+
+    def on_no_clicked(self) -> None:
+        """Handles the no button being clicked."""
+        self.picked_option = False
+        self.close()
+
+    def on_yes_clicked(self) -> None:
+        """Handles the yes button being clicked."""
+        self.picked_option = True
+        self.close()
+
+    def exec_(self):
+        """Executes the yes/no popup. Mimics the behavior of a blocking dialog."""
+        self.show()
+        self.raise_()
+        self.activateWindow()
+
+        loop = QtCore.QEventLoop()
+        self.closed.connect(loop.quit)
+        loop.exec_()
+
+    def closeEvent(self, event):  # noqa: N802
+        """Emits the closed signal when the window is closed."""
+        self.closed.emit()
+        super().closeEvent(event)
+
+
+def display_yes_no_popup(question_text: str) -> bool:
+    """Displays a yes/no popup with the given question text.
+
+    Args:
+        question_text: The question to display.
+
+    Returns:
+        True if the user clicked yes, False if the user clicked no.
+    """
+    yes_no_popup = YesNoPopup(question_text)
+    yes_no_popup.exec_()
+    return yes_no_popup.picked_option
+
+
 class ReceivedMailPopup(NodeMailerWindow):
     """A popup window that display the received mail prompt."""