[safe_mode] Fix setup()-exit return getting trapped in IIFE

safe_mode emits `if (should_enter_safe_mode(...)) return` via
cg.add(RawExpression(...)) to short-circuit the rest of setup() and
boot into safe mode. With setup() split into per-component IIFEs,
that `return` was only exiting the lambda, so the rest of setup() ran
anyway — breaking safe-mode recovery.

Add IIFEUnsafeStatement, a Statement wrapper that marks its containing
component's block for flat emission (no IIFE). safe_mode wraps its
return expression in it. cpp_main_section detects any such statement
in a group and emits that group flat so control-flow constructs like
`return` still affect setup() itself.

IIFEUnsafeStatement.__str__ routes its inner through statement() so
bare Expression subclasses pick up the terminating semicolon. Reported
by @swoboda1337.

esphome/codegen.py

@@ -13,6 +13,7 @@ from esphome.cpp_generator import (  # noqa: F401
     ComponentMarker,
     Expression,
     FlashStringLiteral,
+    IIFEUnsafeStatement,
     LineComment,
     LogStringLiteral,
     MockObj,

esphome/components/safe_mode/__init__.py

@@ -87,7 +87,10 @@ async def to_code(config):
             config[CONF_REBOOT_TIMEOUT],
             config[CONF_BOOT_IS_GOOD_AFTER],
         )
-        cg.add(RawExpression(f"if ({condition}) return"))
+        # Wrap in IIFEUnsafeStatement so cpp_main_section emits this
+        # component's block flat rather than inside an IIFE lambda —
+        # the `return` must exit setup() itself, not just the lambda.
+        cg.add(cg.IIFEUnsafeStatement(RawExpression(f"if ({condition}) return")))
 
     CORE.data[CONF_SAFE_MODE] = {}
     CORE.data[CONF_SAFE_MODE][KEY_PAST_SAFE_MODE] = True

esphome/core/__init__.py

@@ -541,7 +541,7 @@ def _wrap_in_iifes(lines: list[str], max_statements: int) -> list[str]:
     keeping flash small without regressing peak stack."""
     out: list[str] = []
     chunk: list[str] = []
-    depth = 0
+    depth: int = 0
 
     def flush() -> None:
         if not chunk:
@@ -1038,29 +1038,43 @@ class EsphomeCore:
         self.data[KEY_CONTROLLER_REGISTRY_COUNT] = controller_count + 1
 
     @property
-    def cpp_main_section(self):
-        from esphome.cpp_generator import ComponentMarker, statement
+    def cpp_main_section(self) -> str:
+        from esphome.cpp_generator import (
+            ComponentMarker,
+            IIFEUnsafeStatement,
+            statement,
+        )
 
         # Split main_statements at ComponentMarker sentinels and wrap each
         # component's group in an IIFE, sub-splitting at 50 statements so
         # a single heavy component (e.g. a sensor platform with many
         # filter registrations) can't blow the peak chunk frame.
+        # Components that contain an IIFEUnsafeStatement (e.g. safe_mode's
+        # setup-scope `return`) are emitted flat so the statement affects
+        # setup()'s control flow, not the lambda's.
         prefix: list[str] = []
-        components: list[list[str]] = []
-        current = prefix
+        components: list[tuple[list[str], list[bool]]] = []
+        current: list[str] = prefix
+        unsafe_flag: list[bool] = [False]  # unused for prefix
         for exp in self.main_statements:
             if isinstance(exp, ComponentMarker):
                 current = []
-                components.append(current)
+                unsafe_flag = [False]
+                components.append((current, unsafe_flag))
                 continue
+            if isinstance(exp, IIFEUnsafeStatement):
+                unsafe_flag[0] = True
             current.append(str(statement(exp)).rstrip())
 
         if not components:
             return "\n".join(prefix) + "\n\n"
 
-        pieces = list(prefix)
-        for body in components:
-            pieces.extend(_wrap_in_iifes(body, max_statements=50))
+        pieces: list[str] = list(prefix)
+        for body, group_unsafe in components:
+            if group_unsafe[0]:
+                pieces.extend(body)
+            else:
+                pieces.extend(_wrap_in_iifes(body, max_statements=50))
         return "\n".join(pieces) + "\n\n"
 
     @property

esphome/cpp_generator.py

@@ -434,6 +434,25 @@ class LineComment(Statement):
         return "\n".join(parts)
 
 
+class IIFEUnsafeStatement(Statement):
+    """Statement that must not be placed inside an IIFE lambda when
+    ``cpp_main_section`` chunks ``setup()``. Causes the containing
+    component's block to be emitted flat (no IIFE), so constructs that
+    rely on exiting ``setup()`` directly — e.g. safe_mode's
+    ``if (should_enter_safe_mode(...)) return;`` — still work.
+
+    Accepts either a ``Statement`` or a bare ``Expression``; bare
+    expressions are wrapped so they terminate with a semicolon."""
+
+    __slots__ = ("inner",)
+
+    def __init__(self, inner: Expression | Statement) -> None:
+        self.inner = inner
+
+    def __str__(self) -> str:
+        return str(statement(self.inner))
+
+
 class ComponentMarker(Statement):
     """Chunking-boundary sentinel. ``cpp_main_section`` wraps the
     statements between two markers in an IIFE to shorten temporary
@@ -446,10 +465,10 @@ class ComponentMarker(Statement):
 
     __slots__ = ("name",)
 
-    def __init__(self, name: str):
+    def __init__(self, name: str) -> None:
         self.name = name
 
-    def __str__(self):
+    def __str__(self) -> str:
         return f"// component-marker: {self.name}"
 
 

tests/unit_tests/test_core.py

@@ -7,7 +7,12 @@ import pytest
 from strategies import mac_addr_strings
 
 from esphome import const, core
-from esphome.cpp_generator import ComponentMarker, RawStatement
+from esphome.cpp_generator import (
+    ComponentMarker,
+    IIFEUnsafeStatement,
+    RawExpression,
+    RawStatement,
+)
 
 
 class TestHexInt:
@@ -1011,6 +1016,31 @@ def test_cpp_main_section_component_marker_wraps_in_iife() -> None:
     assert "component-marker" not in out
 
 
+def test_cpp_main_section_iife_unsafe_statement_emits_component_flat() -> None:
+    # A component that emits IIFEUnsafeStatement (e.g. safe_mode with
+    # `if (...) return;`) must be emitted flat — a `return` inside an
+    # IIFE would only exit the lambda, not setup().
+    target = core.EsphomeCore()
+    target.main_statements = [
+        ComponentMarker("logger"),
+        RawStatement("new_logger();"),
+        ComponentMarker("safe_mode"),
+        RawStatement("new_safe_mode();"),
+        IIFEUnsafeStatement(RawExpression("if (entering) return")),
+        ComponentMarker("sensor"),
+        RawStatement("new_sensor();"),
+    ]
+    out = target.cpp_main_section
+    # logger and sensor wrapped; safe_mode flat.
+    assert out.count("[]() {") == 2
+    # safe_mode's statements appear at top level, not indented in a lambda.
+    assert "new_safe_mode();" in out
+    assert "if (entering) return;" in out
+    # The IIFEUnsafeStatement wrapper picks up the trailing semicolon
+    # via statement() when inner is a bare Expression.
+    assert "if (entering) return\n" not in out
+
+
 def test_cpp_main_section_comment_only_component_omits_iife() -> None:
     # A component that emits only a ComponentMarker (no statements) adds
     # nothing to the generated output. A neighboring component with