Merge pull request #219 from sparkfun/RC__More_GUI_Updates

Uploader GUI v1.3

Author: PaulZC

Uploader_GUI/RTK_Firmware_Uploader_GUI.py

@@ -13,16 +13,17 @@
 
 Pyinstaller:
 Windows:
-pyinstaller --onefile --clean --noconsole --distpath=./Windows_exe --icon=RTK.ico --add-binary="RTK_Surveyor.ino.partitions.bin;." --add-binary="RTK_Surveyor.ino.bootloader.bin;." --add-binary="boot_app0.bin;." --add-binary="RTK.png;." RTK_Firmware_Uploader_GUI.py
+pyinstaller --onefile --clean --noconsole --distpath=./Windows_exe --icon=RTK.ico --add-binary="RTK_Surveyor_Partitions_4MB.bin;." --add-binary="RTK_Surveyor_Partitions_16MB.bin;." --add-binary="RTK_Surveyor.ino.bootloader.bin;." --add-binary="boot_app0.bin;." --add-binary="RTK.png;." RTK_Firmware_Uploader_GUI.py
 Linux:
-pyinstaller --onefile --clean --noconsole --distpath=./Linux_exe --icon=RTK.ico --add-binary="RTK_Surveyor.ino.partitions.bin:." --add-binary="RTK_Surveyor.ino.bootloader.bin:." --add-binary="boot_app0.bin:." --add-binary="RTK.png:." RTK_Firmware_Uploader_GUI.py
+pyinstaller --onefile --clean --noconsole --distpath=./Linux_exe --icon=RTK.ico --add-binary="RTK_Surveyor_Partitions_4MB.bin:." --add-binary="RTK_Surveyor_Partitions_16MB.bin:." --add-binary="RTK_Surveyor.ino.bootloader.bin:." --add-binary="boot_app0.bin:." --add-binary="RTK.png:." RTK_Firmware_Uploader_GUI.py
 
 Pyinstaller needs:
 RTK_Firmware_Uploader_GUI.py (this file!)
 RTK.ico (icon file for the .exe)
 RTK.png (icon for the GUI widget)
 esptool.py (v3.3, copied from https://github.com/espressif/esptool/releases/tag/v3.3)
-RTK_Surveyor.ino.partitions.bin
+RTK_Surveyor_Partitions_4MB.bin
+RTK_Surveyor_Partitions_16MB.bin
 RTK_Surveyor.ino.bootloader.bin
 boot_app0.bin
 
@@ -35,12 +36,12 @@
 
 from typing import Iterator, Tuple
 
-from PyQt5.QtCore import QSettings, QProcess, QTimer, Qt, QIODevice, pyqtSlot
+from PyQt5.QtCore import QSettings, QProcess, QTimer, Qt, QIODevice, pyqtSlot, QObject
 from PyQt5.QtWidgets import QWidget, QLabel, QComboBox, QGridLayout, \
     QPushButton, QApplication, QLineEdit, QFileDialog, QPlainTextEdit, \
     QAction, QActionGroup, QMenu, QMenuBar, QMainWindow, QMessageBox
 from PyQt5.QtGui import QCloseEvent, QTextCursor, QIcon, QFont
-from PyQt5.QtSerialPort import QSerialPortInfo, QSerialPortInfo
+from PyQt5.QtSerialPort import QSerialPortInfo
 
 import sys
 import os
@@ -55,9 +56,10 @@
 # Setting constants
 SETTING_PORT_NAME = 'port_name'
 SETTING_FILE_LOCATION = 'file_location'
+#SETTING_PARTITION_LOCATION = 'partition_location'
 SETTING_BAUD_RATE = 'baud'
 
-guiVersion = 'v1.2'
+guiVersion = 'v1.3'
 
 def gen_serial_ports() -> Iterator[Tuple[str, str, str]]:
     """Return all available serial ports."""
@@ -70,12 +72,13 @@ def resource_path(relative_path):
     base_path = getattr(sys, '_MEIPASS', os.path.dirname(os.path.abspath(__file__)))
     return os.path.join(base_path, relative_path)
 
-class messageRedirect:
+class messageRedirect(QObject):
     """Wrap a class around a QPlainTextEdit so we can redirect stdout and stderr to it"""
 
-    def __init__(self, edit, out=None) -> None:
+    def __init__(self, edit, out=None, flashSize=None) -> None:
         self.edit = edit
         self.out = out
+        self.flashSize = flashSize
 
     def write(self, msg) -> None:
         if msg.startswith("\r"):
@@ -86,9 +89,21 @@ def write(self, msg) -> None:
             self.edit.insertPlainText(msg)
         self.edit.ensureCursorVisible()
         self.edit.repaint()
+        QApplication.processEvents() # This prevents the circle of doom...
+
         if self.out: # Echo to out (stdout) too if desired
             self.out.write(msg)
 
+        if self.flashSize:
+            if msg.find("Detected flash size: 4MB") >= 0:
+                self.flashSize[0] = 4
+            elif msg.find("Detected flash size: 8MB") >= 0:
+                self.flashSize[0] = 8
+            elif msg.find("Detected flash size: 16MB") >= 0:
+                self.flashSize[0] = 16
+            elif msg.find("Detected flash size: ") >= 0:
+                self.flashSize[0] = 0
+
     def flush(self) -> None:
         None
 
@@ -103,13 +118,12 @@ class MainWidget(QWidget):
     def __init__(self, parent: QWidget = None) -> None:
         super().__init__(parent)
 
-        self.timer=QTimer()
-        self.timer.timeout.connect(self.repaintMessageBox)
+        self.flashSize = [0] # flashSize needs to be mutable. Use a single element list
 
         # File location line edit
-        self.msg_label = QLabel(self.tr('Firmware File:'))
+        self.file_label = QLabel(self.tr('Firmware File:'))
         self.fileLocation_lineedit = QLineEdit()
-        self.msg_label.setBuddy(self.fileLocation_lineedit)
+        self.file_label.setBuddy(self.fileLocation_lineedit)
         self.fileLocation_lineedit.setEnabled(False)
         self.fileLocation_lineedit.returnPressed.connect(self.on_browse_btn_pressed)
 
@@ -118,6 +132,18 @@ def __init__(self, parent: QWidget = None) -> None:
         self.browse_btn.setEnabled(True)
         self.browse_btn.pressed.connect(self.on_browse_btn_pressed)
 
+        # # Partition file location line edit
+        # self.partition_label = QLabel(self.tr('Partition File:'))
+        # self.partitionFileLocation_lineedit = QLineEdit()
+        # self.partition_label.setBuddy(self.partitionFileLocation_lineedit)
+        # self.partitionFileLocation_lineedit.setEnabled(False)
+        # self.partitionFileLocation_lineedit.returnPressed.connect(self.on_partition_browse_btn_pressed)
+
+        # # Browse for new file button
+        # self.partition_browse_btn = QPushButton(self.tr('Browse'))
+        # self.partition_browse_btn.setEnabled(True)
+        # self.partition_browse_btn.pressed.connect(self.on_partition_browse_btn_pressed)
+
         # Port Combobox
         self.port_label = QLabel(self.tr('COM Port:'))
         self.port_combobox = QComboBox()
@@ -152,10 +178,14 @@ def __init__(self, parent: QWidget = None) -> None:
         # Arrange Layout
         layout = QGridLayout()
 
-        layout.addWidget(self.msg_label, 1, 0)
+        layout.addWidget(self.file_label, 1, 0)
         layout.addWidget(self.fileLocation_lineedit, 1, 1)
         layout.addWidget(self.browse_btn, 1, 2)
 
+        # layout.addWidget(self.partition_label, 2, 0)
+        # layout.addWidget(self.partitionFileLocation_lineedit, 2, 1)
+        # layout.addWidget(self.partition_browse_btn, 2, 2)
+
         layout.addWidget(self.port_label, 2, 0)
         layout.addWidget(self.port_combobox, 2, 1)
         layout.addWidget(self.refresh_btn, 2, 2)
@@ -175,20 +205,11 @@ def __init__(self, parent: QWidget = None) -> None:
 
     def writeMessage(self, msg) -> None:
         self.messageBox.moveCursor(QTextCursor.End)
-        self.messageBox.ensureCursorVisible()
+        #self.messageBox.ensureCursorVisible()
         self.messageBox.appendPlainText(msg)
         self.messageBox.ensureCursorVisible()
         self.messageBox.repaint()
-
-    def startTimer(self) -> None:
-        self.timer.start(1000)
-
-    def endTimer(self) -> None:
-        self.timer.stop()
-
-    def repaintMessageBox(self) -> None:
-        self.messageBox.ensureCursorVisible()
-        self.messageBox.repaint()
+        QApplication.processEvents()
 
     def _load_settings(self) -> None:
         """Load settings on startup."""
@@ -202,6 +223,12 @@ def _load_settings(self) -> None:
         if lastFile is not None:
             self.fileLocation_lineedit.setText(lastFile)
 
+        # lastFile = self.settings.value(SETTING_PARTITION_LOCATION)
+        # if lastFile is not None:
+        #     self.partitionFileLocation_lineedit.setText(lastFile)
+        # else:
+        #     self.partitionFileLocation_lineedit.setText(resource_path("RTK_Surveyor.ino.partitions.bin"))
+
         baud = self.settings.value(SETTING_BAUD_RATE)
         if baud is not None:
             index = self.baud_combobox.findData(baud)
@@ -212,6 +239,7 @@ def _save_settings(self) -> None:
         """Save settings on shutdown."""
         self.settings.setValue(SETTING_PORT_NAME, self.port)
         self.settings.setValue(SETTING_FILE_LOCATION, self.theFileName)
+        # self.settings.setValue(SETTING_PARTITION_LOCATION, self.thePartitionFileName)
         self.settings.setValue(SETTING_BAUD_RATE, self.baudRate)
 
     def _clean_settings(self) -> None:
@@ -264,15 +292,18 @@ def theFileName(self) -> str:
         """Return the file name."""
         return self.fileLocation_lineedit.text()
 
+    # @property
+    # def thePartitionFileName(self) -> str:
+    #     """Return the partition file name."""
+    #     return self.partitionFileLocation_lineedit.text()
+
     def closeEvent(self, event: QCloseEvent) -> None:
         """Handle Close event of the Widget."""
         try:
             self._save_settings()
         except:
             pass
 
-        self.endTimer()
-        
         event.accept()
 
     def on_refresh_btn_pressed(self) -> None:
@@ -291,6 +322,18 @@ def on_browse_btn_pressed(self) -> None:
         if fileName:
             self.fileLocation_lineedit.setText(fileName)
 
+    # def on_partition_browse_btn_pressed(self) -> None:
+    #     """Open dialog to select partition bin file."""
+    #     options = QFileDialog.Options()
+    #     fileName, _ = QFileDialog.getOpenFileName(
+    #         None,
+    #         "Select Partition File",
+    #         "",
+    #         "Parition Files (*.bin);;All Files (*)",
+    #         options=options)
+    #     if fileName:
+    #         self.partitionFileLocation_lineedit.setText(fileName)
+
     def on_upload_btn_pressed(self) -> None:
         """Upload the firmware"""
         portAvailable = False
@@ -303,7 +346,7 @@ def on_upload_btn_pressed(self) -> None:
 
         fileExists = False
         try:
-            f = open(self.fileLocation_lineedit.text())
+            f = open(self.theFileName)
             fileExists = True
         except IOError:
             fileExists = False
@@ -313,11 +356,67 @@ def on_upload_btn_pressed(self) -> None:
                 return
             f.close()
 
+        # fileExists = False
+        # try:
+        #     f = open(self.thePartitionFileName)
+        #     fileExists = True
+        # except IOError:
+        #     fileExists = False
+        # finally:
+        #     if (fileExists == False):
+        #         self.writeMessage("File Not Found")
+        #         return
+        #     f.close()
+
         try:
             self._save_settings() # Save the settings in case the upload crashes
         except:
             pass
 
+        self.flashSize[0] = 0
+
+        self.writeMessage("Detecting flash size\n\n")
+
+        command = []
+        command.extend(["--chip","esp32"])
+        command.extend(["--port",self.port])
+        command.extend(["--baud",self.baudRate])
+        command.extend(["flash_id"])
+
+        try:
+            esptool.main(command)
+        except (ValueError, IOError, FatalError, ImportError, NotImplementedInROMError, UnsupportedCommandError, NotSupportedError, RuntimeError) as err:
+            self.writeMessage(str(err))
+            self.messageBox.ensureCursorVisible()
+            self.messageBox.repaint()
+            return
+        except:
+            self.messageBox.ensureCursorVisible()
+            self.messageBox.repaint()
+            return
+
+        if self.flashSize[0] == 0:
+            self.writeMessage("Flash size not detected! Defaulting to 16MB\n")
+            self.flashSize[0] = 16
+        else:
+            self.writeMessage("Flash size is " + str(self.flashSize[0]) + "MB\n")
+
+        thePartitionFileName = ''
+        firmwareSizeCorrect = True
+        if self.flashSize[0] == 16:
+            thePartitionFileName = resource_path("RTK_Surveyor_Partitions_16MB.bin")
+            # if self.theFileName.find("16MB") < 0:
+            #     firmwareSizeCorrect = False
+        else:
+            thePartitionFileName = resource_path("RTK_Surveyor_Partitions_4MB.bin")
+            # if self.theFileName.find("4MB") < 0:
+            #     firmwareSizeCorrect = False
+
+        if firmwareSizeCorrect == False:
+            reply = QMessageBox.warning(self, "Firmware size mismatch", "Do you want to continue?", QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
+            if reply == QMessageBox.No:
+                return
+
         self.writeMessage("Uploading firmware\n")
 
         command = []
@@ -327,22 +426,23 @@ def on_upload_btn_pressed(self) -> None:
         command.extend(["--baud",self.baudRate])
         command.extend(["--before","default_reset","--after","hard_reset","write_flash","-z","--flash_mode","dio","--flash_freq","80m","--flash_size","detect"])
         command.extend(["0x1000",resource_path("RTK_Surveyor.ino.bootloader.bin")])
-        command.extend(["0x8000",resource_path("RTK_Surveyor.ino.partitions.bin")])
+        command.extend(["0x8000",thePartitionFileName])
         command.extend(["0xe000",resource_path("boot_app0.bin")])
         command.extend(["0x10000",self.theFileName])
 
         self.writeMessage("Command: esptool.main(%s)\n\n" % " ".join(command))
 
         #print("python esptool.py %s\n\n" % " ".join(command)) # Useful for debugging - cut and paste into a command prompt
 
-        self.startTimer()
-
         try:
             esptool.main(command)
         except (ValueError, IOError, FatalError, ImportError, NotImplementedInROMError, UnsupportedCommandError, NotSupportedError, RuntimeError) as err:
             self.writeMessage(str(err))
+        except:
+            pass
 
-        self.endTimer()
+        self.messageBox.ensureCursorVisible()
+        self.messageBox.repaint()
 
 if __name__ == '__main__':
     from sys import exit as sysExit
@@ -352,10 +452,10 @@ def on_upload_btn_pressed(self) -> None:
     app.setWindowIcon(QIcon(resource_path("RTK.png")))
     w = MainWidget()
     if 1: # Change to 0 to have the messages echoed on stdout
-        sys.stdout = messageRedirect(w.messageBox) # Divert stdout to messageBox
+        sys.stdout = messageRedirect(w.messageBox, flashSize=w.flashSize) # Divert stdout to messageBox. Report flash size via flashSize
         sys.stderr = messageRedirect(w.messageBox) # Divert stderr to messageBox
     else:
-        sys.stdout = messageRedirect(w.messageBox, sys.stdout) # Echo to stdout too
-        sys.stderr = messageRedirect(w.messageBox, sys.stderr) # Echo to stderr too
+        sys.stdout = messageRedirect(w.messageBox, flashSize=w.flashSize, out=sys.stdout) # Echo to stdout too
+        sys.stderr = messageRedirect(w.messageBox, out=sys.stderr) # Echo to stderr too
     w.show()
     sys.exit(app.exec_())