feat(cmds): Encapsulate logic for running the stub flasher in run_stub

docs/en/esptool/scripting.rst

@@ -41,20 +41,20 @@ This example demonstrates writing two binary files using high-level commands:
 
 .. code-block:: python
 
-    from esptool.cmds import detect_chip, attach_flash, reset_chip, write_flash
+    from esptool.cmds import detect_chip, attach_flash, reset_chip, run_stub, write_flash
 
     PORT = "/dev/ttyACM0"
     BOOTLOADER = "bootloader.bin"
     FIRMWARE = "firmware.bin"
 
     with detect_chip(PORT) as esp:
-        esp = esp.run_stub()  # Skip this line to avoid running the stub flasher
+        esp = run_stub(esp)  # Skip this line to avoid running the stub flasher
         attach_flash(esp)  # Attach the flash memory chip, required for flash operations
         with open(BOOTLOADER, "rb") as bl_file, open(FIRMWARE, "rb") as fw_file:
             write_flash(esp, [(0, bl_file), (0x1000, fw_file)])  # Write the binary files
         reset_chip(esp, "hard_reset")  # Reset the chip
 
-- The ``esp`` object has to be replaced with the stub flasher object returned by ``esp.run_stub()`` when the stub flasher is activated. This step can be skipped if the stub flasher is not needed.
+- The ``esp`` object has to be replaced with the stub flasher object returned by ``run_stub(esp)`` when the stub flasher is activated. This step can be skipped if the stub flasher is not needed.
 - Running ``attach_flash(esp)`` is required for any flash-memory-related operations to work.
 - Using the ``esp`` object in a context manager ensures the port gets closed properly after the block is executed.
 
@@ -70,6 +70,7 @@ The following example demonstrates running a series of flash memory operations i
     flash_id,
     read_flash,
     reset_chip,
+    run_stub,
     verify_flash,
     write_flash,
     )
@@ -81,7 +82,7 @@ The following example demonstrates running a series of flash memory operations i
 
     with ESP32ROM(PORT) as esp:
         esp.connect()  # Connect to the ESP chip, needed when ESP32ROM is instantiated directly
-        esp = esp.run_stub()  # Run the stub loader (optional)
+        esp = run_stub(esp)  # Run the stub loader (optional)
         attach_flash(esp)  # Attach the flash memory chip, required for flash operations
         flash_id(esp)  # Print information about the flash chip
         erase_flash(esp)  # Erase the flash memory first
@@ -103,6 +104,8 @@ Chip Control Operations
 
 .. autofunction:: esptool.cmds.detect_chip
 
+.. autofunction:: esptool.cmds.run_stub
+
 .. autofunction:: esptool.cmds.load_ram
 
 .. autofunction:: esptool.cmds.run

esptool/__init__.py

@@ -24,6 +24,7 @@ __all__ = [
     "read_mem",
     "reset_chip",
     "run",
+    "run_stub",
     "verify_flash",
     "version",
     "write_flash",
@@ -62,6 +63,7 @@ from esptool.cmds import (
     read_mem,
     reset_chip,
     run,
+    run_stub,
     verify_flash,
     version,
     write_flash,
@@ -449,36 +451,7 @@ def prepare_esp_object(ctx):
     ############################
 
     if not ctx.obj["no_stub"]:
-        if esp.secure_download_mode:
-            log.warning(
-                "Stub loader is not supported in Secure Download Mode, "
-                "it has been disabled. Set --no-stub to suppress this warning."
-            )
-        elif not esp.IS_STUB and esp.stub_is_disabled:
-            log.warning(
-                "Stub loader has been disabled for compatibility, "
-                "set --no-stub to suppress this warning."
-            )
-        elif esp.CHIP_NAME in [
-            "ESP32-H21",
-            "ESP32-H4",
-        ]:  # TODO: [ESP32H21] IDF-11509   [ESP32H4] IDF-12271
-            log.warning(
-                f"Stub loader is not yet supported on {esp.CHIP_NAME}, "
-                "it has been disabled. Set --no-stub to suppress this warning."
-            )
-        else:
-            try:
-                esp = esp.run_stub()
-            except Exception:
-                # The CH9102 bridge (PID: 0x55D4) can have issues on MacOS
-                if sys.platform == "darwin" and esp._get_pid() == 0x55D4:
-                    log.print()
-                    log.note(
-                        "If issues persist, "
-                        "try installing the WCH USB-to-Serial MacOS driver."
-                    )
-                raise
+        esp = run_stub(esp)
 
     # 4) Configure the baud rate and voltage regulator
     ##################################################

esptool/cmds.py

@@ -7,6 +7,7 @@ import hashlib
 import io
 import os
 import struct
+import sys
 import time
 import zlib
 import itertools
@@ -1547,6 +1548,52 @@ def reset_chip(esp: ESPLoader, reset_mode: str = "hard_reset") -> None:
         raise FatalError(f"Invalid reset mode: {reset_mode}")
 
 
+def run_stub(esp: ESPLoader) -> ESPLoader:
+    """
+    Load and execute the stub loader on the ESP device. If stub loading
+    is not supported or is explicitly disabled, warnings are logged.
+
+    Args:
+        esp: Initiated esp object connected to a real device.
+
+    Returns:
+        ESPLoader: The esp instance, either as a stub child class in a state
+        where the stub has been executed, or in its original state
+        if the stub loader is disabled or unsupported.
+    """
+    if esp.secure_download_mode:
+        log.warning(
+            "Stub loader is not supported in Secure Download Mode, "
+            "it has been disabled. Set --no-stub to suppress this warning."
+        )
+    elif not esp.IS_STUB and esp.stub_is_disabled:
+        log.warning(
+            "Stub loader has been disabled for compatibility, "
+            "set --no-stub to suppress this warning."
+        )
+    elif esp.CHIP_NAME in [
+        "ESP32-H21",
+        "ESP32-H4",
+    ]:  # TODO: [ESP32H21] IDF-11509   [ESP32H4] IDF-12271
+        log.warning(
+            f"Stub loader is not yet supported on {esp.CHIP_NAME}, "
+            "it has been disabled. Set --no-stub to suppress this warning."
+        )
+    else:
+        try:
+            return esp.run_stub()
+        except Exception:
+            # The CH9102 bridge (PID: 0x55D4) can have issues on MacOS
+            if sys.platform == "darwin" and esp._get_pid() == 0x55D4:
+                log.print()
+                log.note(
+                    "If issues persist, "
+                    "try installing the WCH USB-to-Serial MacOS driver."
+                )
+            raise
+    return esp
+
+
 # Commands that don't require an ESP object (image manipulation, etc.)
 ######################################################################
 

esptool/loader.py

@@ -1162,13 +1162,13 @@ class ESPLoader(object):
                 % (arg, ", ".join(cls.FLASH_FREQUENCY.keys()))
             )
 
-    def run_stub(self, stub=None):
+    def run_stub(self, stub: StubFlasher | None = None) -> "ESPLoader":
         if stub is None:
             stub = StubFlasher(self.CHIP_NAME)
 
         if self.sync_stub_detected:
             log.print("Stub is already running. No upload is necessary.")
-            return self.STUB_CLASS(self)
+            return self.STUB_CLASS(self) if self.STUB_CLASS is not None else self
 
         # Upload
         log.print("Uploading stub...")
@@ -1196,7 +1196,7 @@ class ESPLoader(object):
         if p != b"OHAI":
             raise FatalError(f"Failed to start stub. Unexpected response: {p}")
         log.print("Stub running...")
-        return self.STUB_CLASS(self)
+        return self.STUB_CLASS(self) if self.STUB_CLASS is not None else self
 
     @stub_and_esp32_function_only
     def flash_defl_begin(self, size, compsize, offset):