Fix showing progress

1 files changed, 14 insertions, 4 deletions

diff --git a/plugins/UltimakerMachineActions/UpgradeFirmwareMachineAction.py b/plugins/UltimakerMachineActions/UpgradeFirmwareMachineAction.py
index 671ed22d5a..478ea9b6bb 100644
--- a/plugins/UltimakerMachineActions/UpgradeFirmwareMachineAction.py
+++ b/plugins/UltimakerMachineActions/UpgradeFirmwareMachineAction.py
@@ -6,6 +6,7 @@ from UM.Settings.DefinitionContainer import DefinitionContainer
 from cura.MachineAction import MachineAction
 from UM.i18n import i18nCatalog
 from UM.Settings.ContainerRegistry import ContainerRegistry
+from cura.PrinterOutput.FirmwareUpdater import FirmwareUpdateState
 
 from PyQt5.QtCore import pyqtSignal, pyqtProperty, QObject
 from typing import Optional
@@ -13,6 +14,7 @@ from typing import Optional
 MYPY = False
 if MYPY:
     from cura.PrinterOutput.FirmwareUpdater import FirmwareUpdater
+    from cura.PrinterOutput.PrinterOutputDevice import PrinterOutputDevice
 
 catalog = i18nCatalog("cura")
 
@@ -23,7 +25,8 @@ class UpgradeFirmwareMachineAction(MachineAction):
         self._qml_url = "UpgradeFirmwareMachineAction.qml"
         ContainerRegistry.getInstance().containerAdded.connect(self._onContainerAdded)
 
-        self._active_output_device = None
+        self._active_output_device = None #type: Optional[PrinterOutputDevice]
+        self._active_firmware_updater = None #type: Optional[FirmwareUpdater]
 
         Application.getInstance().engineCreatedSignal.connect(self._onEngineCreated)
 
@@ -38,9 +41,10 @@ class UpgradeFirmwareMachineAction(MachineAction):
     def _onOutputDevicesChanged(self) -> None:
         if self._active_output_device:
             self._active_output_device.activePrinter.getController().canUpdateFirmwareChanged.disconnect(self._onControllerCanUpdateFirmwareChanged)
+
         output_devices = Application.getInstance().getMachineManager().printerOutputDevices
-        print(output_devices)
         self._active_output_device = output_devices[0] if output_devices else None
+
         if self._active_output_device:
             self._active_output_device.activePrinter.getController().canUpdateFirmwareChanged.connect(self._onControllerCanUpdateFirmwareChanged)
 
@@ -53,6 +57,12 @@ class UpgradeFirmwareMachineAction(MachineAction):
     @pyqtProperty(QObject, notify = outputDeviceCanUpdateFirmwareChanged)
     def firmwareUpdater(self) -> Optional["firmwareUpdater"]:
         if self._active_output_device and self._active_output_device.activePrinter.getController().can_update_firmware:
-            return self._active_output_device.getFirmwareUpdater()
+            self._active_firmware_updater = self._active_output_device.getFirmwareUpdater()
+            return self._active_firmware_updater
+
+        elif self._active_firmware_updater and self._active_firmware_updater.firmwareUpdateState not in [FirmwareUpdateState.idle, FirmwareUpdateState.completed]:
+            # During a firmware update, the PrinterOutputDevice is disconnected but the FirmwareUpdater is still there
+            return self._active_firmware_updater
 
-        return None
\ No newline at end of file
+        self._active_firmware_updater = None
+        return None