esphome-lsm6ds/esphome/components/esp32/post_build.py.script

Import("env")

import os
import json
import shutil
import pathlib
import itertools

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:
            with flasher_args_path.open() as f:
                flash_data = json.load(f)
            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:
                    print(f"Info: {file_path.name} not found - skipping")
        except Exception as e:
            print(f"Warning: Failed to parse flasher_args.json - {e}")

    # 2. Try FLASH_EXTRA_IMAGES if flasher_args.json failed or was empty
    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")

    # 3. Final fallback: guess standard image locations
    if not sections:
        print("Fallback: guessing legacy image paths")
        guesses = [
            ("0x0", build_dir / "bootloader" / "bootloader.bin"),
            ("0x8000", build_dir / "partition_table" / "partition-table.bin"),
            ("0xe000", build_dir / "ota_data_initial.bin"),
            ("0x10000", firmware_path)
        ]
        for addr, file_path in guesses:
            if file_path.exists():
                sections.append((addr, str(file_path)))
            else:
                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
env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", merge_factory_bin)
env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", esp32_copy_ota_bin)