Update post_build.py.script to Fix #7137 (#9578)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>

esphome/components/esp32/post_build.py.script

@@ -1,77 +1,112 @@
-# Source https://github.com/letscontrolit/ESPEasy/pull/3845#issuecomment-1005864664
+Import("env")
-
-# pylint: disable=E0602
-Import("env")  # noqa
 
 import os
+import json
 import shutil
+import pathlib
+import itertools
 
-if os.environ.get("ESPHOME_USE_SUBPROCESS") is None:
+def merge_factory_bin(source, target, env):
+    """
+    Merges all flash sections into a single .factory.bin using esptool.
+    Attempts multiple methods to detect image layout: flasher_args.json, FLASH_EXTRA_IMAGES, fallback guesses.
+    """
+    firmware_name = os.path.basename(env.subst("$PROGNAME")) + ".bin"
+    build_dir = pathlib.Path(env.subst("$BUILD_DIR"))
+    firmware_path = build_dir / firmware_name
+    flash_size = env.BoardConfig().get("upload.flash_size", "4MB")
+    chip = env.BoardConfig().get("build.mcu", "esp32")
+
+    sections = []
+    flasher_args_path = build_dir / "flasher_args.json"
+
+    # 1. Try flasher_args.json
+    if flasher_args_path.exists():
         try:
-        import esptool
+            with flasher_args_path.open() as f:
-    except ImportError:
+                flash_data = json.load(f)
-        env.Execute("$PYTHONEXE -m pip install esptool")
+            for addr, fname in sorted(flash_data["flash_files"].items(), key=lambda kv: int(kv[0], 16)):
+                file_path = pathlib.Path(fname)
+                if file_path.exists():
+                    sections.append((addr, str(file_path)))
                 else:
-    import subprocess
+                    print(f"Info: {file_path.name} not found - skipping")
-from SCons.Script import ARGUMENTS
+        except Exception as e:
+            print(f"Warning: Failed to parse flasher_args.json - {e}")
 
-# Copy over the default sdkconfig.
+    # 2. Try FLASH_EXTRA_IMAGES if flasher_args.json failed or was empty
-from os import path
+    if not sections:
+        flash_images = env.get("FLASH_EXTRA_IMAGES")
+        if flash_images:
+            print("Using FLASH_EXTRA_IMAGES from PlatformIO environment")
+            # flatten any nested lists
+            flat = list(itertools.chain.from_iterable(
+                x if isinstance(x, (list, tuple)) else [x] for x in flash_images
+            ))
+            entries = [env.subst(x) for x in flat]
+            for i in range(0, len(entries) - 1, 2):
+                addr, fname = entries[i], entries[i + 1]
+                if isinstance(fname, (list, tuple)):
+                    print(f"Warning: Skipping malformed FLASH_EXTRA_IMAGES entry: {fname}")
+                    continue
+                file_path = pathlib.Path(str(fname))
+                if file_path.exists():
+                    sections.append((addr, str(file_path)))
+                else:
+                    print(f"Info: {file_path.name} not found — skipping")
 
-if path.exists("./sdkconfig.defaults"):
+    # 3. Final fallback: guess standard image locations
-    os.makedirs(".temp", exist_ok=True)
+    if not sections:
-    shutil.copy("./sdkconfig.defaults", "./.temp/sdkconfig-esp32-idf")
+        print("Fallback: guessing legacy image paths")
-
+        guesses = [
-
+            ("0x0", build_dir / "bootloader" / "bootloader.bin"),
-def esp32_create_combined_bin(source, target, env):
+            ("0x8000", build_dir / "partition_table" / "partition-table.bin"),
-    verbose = bool(int(ARGUMENTS.get("PIOVERBOSE", "0")))
+            ("0xe000", build_dir / "ota_data_initial.bin"),
-    if verbose:
+            ("0x10000", firmware_path)
-        print("Generating combined binary for serial flashing")
-    app_offset = 0x10000
-
-    new_file_name = env.subst("$BUILD_DIR/${PROGNAME}.factory.bin")
-    sections = env.subst(env.get("FLASH_EXTRA_IMAGES"))
-    firmware_name = env.subst("$BUILD_DIR/${PROGNAME}.bin")
-    chip = env.get("BOARD_MCU")
-    flash_size = env.BoardConfig().get("upload.flash_size")
-    cmd = [
-        "--chip",
-        chip,
-        "merge_bin",
-        "-o",
-        new_file_name,
-        "--flash_size",
-        flash_size,
         ]
-    if verbose:
+        for addr, file_path in guesses:
-        print("    Offset | File")
+            if file_path.exists():
-    for section in sections:
+                sections.append((addr, str(file_path)))
-        sect_adr, sect_file = section.split(" ", 1)
-        if verbose:
-            print(f" -  {sect_adr} | {sect_file}")
-        cmd += [sect_adr, sect_file]
-
-    cmd += [hex(app_offset), firmware_name]
-
-    if verbose:
-        print(f" - {hex(app_offset)} | {firmware_name}")
-        print()
-        print(f"Using esptool.py arguments: {' '.join(cmd)}")
-        print()
-
-    if os.environ.get("ESPHOME_USE_SUBPROCESS") is None:
-        esptool.main(cmd)
             else:
-        subprocess.run(["esptool.py", *cmd])
+                print(f"Info: {file_path.name} not found — skipping")
 
+    # If no valid sections found, skip merge
+    if not sections:
+        print("No valid flash sections found — skipping .factory.bin creation.")
+        return
+
+    output_path = firmware_path.with_suffix(".factory.bin")
+    cmd = [
+        "--chip", chip,
+        "merge_bin",
+        "--flash_size", flash_size,
+        "--output", str(output_path)
+    ]
+    for addr, file_path in sections:
+        cmd += [addr, file_path]
+
+    print(f"Merging binaries into {output_path}")
+    result = env.Execute(
+        env.VerboseAction(
+            f"{env.subst('$PYTHONEXE')} -m esptool " + " ".join(cmd),
+            "Merging binaries with esptool"
+        )
+    )
+
+    if result == 0:
+        print(f"Successfully created {output_path}")
+    else:
+        print(f"Error: esptool merge_bin failed with code {result}")
 
 def esp32_copy_ota_bin(source, target, env):
+    """
+    Copy the main firmware to a .ota.bin file for compatibility with ESPHome OTA tools.
+    """
     firmware_name = env.subst("$BUILD_DIR/${PROGNAME}.bin")
     new_file_name = env.subst("$BUILD_DIR/${PROGNAME}.ota.bin")
-
     shutil.copyfile(firmware_name, new_file_name)
+    print(f"Copied firmware to {new_file_name}")
 
-
+# Run merge first, then ota copy second
-# pylint: disable=E0602
+env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", merge_factory_bin)
-env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", esp32_create_combined_bin)  # noqa
+env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", esp32_copy_ota_bin)
-env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", esp32_copy_ota_bin)  # noqa