From 5f2582efcd21a99941a1d05d1772b5f85c10628b Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Fri, 17 Apr 2026 18:12:14 -0500 Subject: [PATCH] [safe_mode] Fix setup()-exit return getting trapped in IIFE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 1 + esphome/components/safe_mode/__init__.py | 5 +++- esphome/core/__init__.py | 32 +++++++++++++++++------- esphome/cpp_generator.py | 23 +++++++++++++++-- tests/unit_tests/test_core.py | 32 +++++++++++++++++++++++- 5 files changed, 80 insertions(+), 13 deletions(-) diff --git a/esphome/codegen.py b/esphome/codegen.py index ffaf368bd45..64ce0da47ed 100644 --- a/esphome/codegen.py +++ b/esphome/codegen.py @@ -13,6 +13,7 @@ from esphome.cpp_generator import ( # noqa: F401 ComponentMarker, Expression, FlashStringLiteral, + IIFEUnsafeStatement, LineComment, LogStringLiteral, MockObj, diff --git a/esphome/components/safe_mode/__init__.py b/esphome/components/safe_mode/__init__.py index 6df0ba78b1f..78c757d293f 100644 --- a/esphome/components/safe_mode/__init__.py +++ b/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 diff --git a/esphome/core/__init__.py b/esphome/core/__init__.py index 8bc4e6e1534..3423b3d953f 100644 --- a/esphome/core/__init__.py +++ b/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 diff --git a/esphome/cpp_generator.py b/esphome/cpp_generator.py index a9dca970b64..c0c4ad5ed66 100644 --- a/esphome/cpp_generator.py +++ b/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}" diff --git a/tests/unit_tests/test_core.py b/tests/unit_tests/test_core.py index e4a52a78080..96190a538ad 100644 --- a/tests/unit_tests/test_core.py +++ b/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