Add cg.with_local_variable (#3577)

esphome/codegen.py

@@ -22,6 +22,7 @@ from esphome.cpp_generator import (  # noqa
     static_const_array,
     statement,
     variable,
+    with_local_variable,
     new_variable,
     Pvariable,
     new_Pvariable,

esphome/components/wifi/__init__.py

@@ -332,8 +332,7 @@ def manual_ip(config):
     )
 
 
-def wifi_network(config, static_ip):
-    ap = cg.variable(config[CONF_ID], WiFiAP())
+def wifi_network(config, ap, static_ip):
     if CONF_SSID in config:
         cg.add(ap.set_ssid(config[CONF_SSID]))
     if CONF_PASSWORD in config:
@@ -360,14 +359,21 @@ async def to_code(config):
     var = cg.new_Pvariable(config[CONF_ID])
     cg.add(var.set_use_address(config[CONF_USE_ADDRESS]))
 
-    for network in config.get(CONF_NETWORKS, []):
+    def add_sta(ap, network):
         ip_config = network.get(CONF_MANUAL_IP, config.get(CONF_MANUAL_IP))
-        cg.add(var.add_sta(wifi_network(network, ip_config)))
+        cg.add(var.add_sta(wifi_network(network, ap, ip_config)))
+
+    for network in config.get(CONF_NETWORKS, []):
+        cg.with_local_variable(network[CONF_ID], WiFiAP(), add_sta, network)
 
     if CONF_AP in config:
         conf = config[CONF_AP]
         ip_config = conf.get(CONF_MANUAL_IP, config.get(CONF_MANUAL_IP))
-        cg.add(var.set_ap(wifi_network(conf, ip_config)))
+        cg.with_local_variable(
+            conf[CONF_ID],
+            WiFiAP(),
+            lambda ap: cg.add(var.set_ap(wifi_network(conf, ap, ip_config))),
+        )
         cg.add(var.set_ap_timeout(conf[CONF_AP_TIMEOUT]))
 
     cg.add(var.set_reboot_timeout(config[CONF_REBOOT_TIMEOUT]))

esphome/cpp_generator.py

@@ -5,7 +5,17 @@ import re
 from esphome.yaml_util import ESPHomeDataBase
 
 # pylint: disable=unused-import, wrong-import-order
-from typing import Any, Generator, List, Optional, Tuple, Type, Union, Sequence
+from typing import (
+    Any,
+    Callable,
+    Generator,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+    Sequence,
+)
 
 from esphome.core import (  # noqa
     CORE,
@@ -468,7 +478,9 @@ def statement(expression: Union[Expression, Statement]) -> Statement:
     return ExpressionStatement(expression)
 
 
-def variable(id_: ID, rhs: SafeExpType, type_: "MockObj" = None) -> "MockObj":
+def variable(
+    id_: ID, rhs: SafeExpType, type_: "MockObj" = None, register=True
+) -> "MockObj":
     """Declare a new variable, not pointer type, in the code generation.
 
     :param id_: The ID used to declare the variable.
@@ -485,10 +497,37 @@ def variable(id_: ID, rhs: SafeExpType, type_: "MockObj" = None) -> "MockObj":
         id_.type = type_
     assignment = AssignmentExpression(id_.type, "", id_, rhs)
     CORE.add(assignment)
+    if register:
         CORE.register_variable(id_, obj)
     return obj
 
 
+def with_local_variable(
+    id_: ID, rhs: SafeExpType, callback: Callable[["MockObj"], None], *args
+) -> None:
+    """Declare a new variable, not pointer type, in the code generation, within a scoped block
+    The variable is only usable within the callback
+    The callback cannot be async.
+
+    :param id_: The ID used to declare the variable.
+    :param rhs: The expression to place on the right hand side of the assignment.
+    :param callback: The function to invoke that will receive the temporary variable
+    :param args: args to pass to the callback in addition to the temporary variable
+
+    """
+
+    # throw if the callback is async:
+    assert not inspect.iscoroutinefunction(
+        callback
+    ), "with_local_variable() callback cannot be async!"
+
+    CORE.add(RawStatement("{"))  # output opening curly brace
+    obj = variable(id_, rhs, None, True)
+    # invoke user-provided callback to generate code with this local variable
+    callback(obj, *args)
+    CORE.add(RawStatement("}"))  # output closing curly brace
+
+
 def new_variable(id_: ID, rhs: SafeExpType, type_: "MockObj" = None) -> "MockObj":
     """Declare and define a new variable, not pointer type, in the code generation.