From 0e9c85ed2a5c80a51a550efea1b99156bbaacc63 Mon Sep 17 00:00:00 2001
From: "radim.karnis" <radim.karnis@espressif.com>
Date: Thu, 30 Mar 2023 14:52:45 +0200
Subject: [PATCH] fix: Set flash parameters even with --flash_size keep

Related to https://github.com/espressif/esp-idf/issues/10788

Related to https://github.com/espressif/esp-idf/issues/10959
---
 esptool/__init__.py | 14 ++++++++++----
 esptool/cmds.py     | 34 +++++++++++++++++++++-------------
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/esptool/__init__.py b/esptool/__init__.py
index 88a021813..e2cebff88 100644
--- a/esptool/__init__.py
+++ b/esptool/__init__.py
@@ -811,11 +811,17 @@ def flash_xmc_startup():
 
         if hasattr(args, "flash_size"):
             print("Configuring flash size...")
-            detect_flash_size(esp, args)
-            if args.flash_size != "keep":  # TODO: should set this even with 'keep'
-                esp.flash_set_parameters(flash_size_bytes(args.flash_size))
+            if args.flash_size == "detect":
+                flash_size = detect_flash_size(esp, args)
+            elif args.flash_size == "keep":
+                flash_size = detect_flash_size(esp, args=None)
+            else:
+                flash_size = args.flash_size
+
+            if flash_size is not None:  # Secure download mode
+                esp.flash_set_parameters(flash_size_bytes(flash_size))
                 # Check if stub supports chosen flash size
-                if esp.IS_STUB and args.flash_size in ("32MB", "64MB", "128MB"):
+                if esp.IS_STUB and flash_size in ("32MB", "64MB", "128MB"):
                     print(
                         "WARNING: Flasher stub doesn't fully support flash size larger "
                         "than 16MB, in case of failure use --no-stub."
diff --git a/esptool/cmds.py b/esptool/cmds.py
index c02fbec5e..fc3825c84 100644
--- a/esptool/cmds.py
+++ b/esptool/cmds.py
@@ -208,24 +208,30 @@ def dump_mem(esp, args):
     print("Done!")
 
 
-def detect_flash_size(esp, args):
-    if args.flash_size == "detect":
-        if esp.secure_download_mode:
+def detect_flash_size(esp, args=None):
+    # TODO: Remove the dependency on args in the next major release (v5.0)
+    if esp.secure_download_mode:
+        if args is not None and args.flash_size == "detect":
             raise FatalError(
                 "Detecting flash size is not supported in secure download mode. "
                 "Need to manually specify flash size."
             )
-        flash_id = esp.flash_id()
-        size_id = flash_id >> 16
-        args.flash_size = DETECTED_FLASH_SIZES.get(size_id)
-        if args.flash_size is None:
+        else:
+            return None
+    flash_id = esp.flash_id()
+    size_id = flash_id >> 16
+    flash_size = DETECTED_FLASH_SIZES.get(size_id)
+    if args is not None and args.flash_size == "detect":
+        if flash_size is None:
+            flash_size = "4MB"
             print(
-                "Warning: Could not auto-detect Flash size (FlashID=0x%x, SizeID=0x%x),"
-                " defaulting to 4MB" % (flash_id, size_id)
+                "Warning: Could not auto-detect Flash size "
+                f"(FlashID={flash_id:#x}, SizeID={size_id:#x}), defaulting to 4MB"
             )
-            args.flash_size = "4MB"
         else:
-            print("Auto-detected Flash size:", args.flash_size)
+            print("Auto-detected Flash size:", flash_size)
+        args.flash_size = flash_size
+    return flash_size
 
 
 def _update_image_flash_params(esp, address, args, image):
@@ -456,8 +462,10 @@ def write_flash(esp, args):
                 )
 
     # verify file sizes fit in flash
-    if args.flash_size != "keep":  # TODO: check this even with 'keep'
-        flash_end = flash_size_bytes(args.flash_size)
+    flash_end = flash_size_bytes(
+        detect_flash_size(esp) if args.flash_size == "keep" else args.flash_size
+    )
+    if flash_end is not None:  # Secure download mode
         for address, argfile in args.addr_filename:
             argfile.seek(0, os.SEEK_END)
             if address + argfile.tell() > flash_end: