[substitutions] !extend and !remove now support substitutions and jinja (#11203)

esphome/config.py

@@ -12,7 +12,7 @@ from typing import Any
 import voluptuous as vol
 
 from esphome import core, loader, pins, yaml_util
-from esphome.config_helpers import Extend, Remove, merge_dicts_ordered
+from esphome.config_helpers import Extend, Remove, merge_config, merge_dicts_ordered
 import esphome.config_validation as cv
 from esphome.const import (
     CONF_ESPHOME,
@@ -324,13 +324,7 @@ def iter_ids(config, path=None):
             yield from iter_ids(value, path + [key])
 
 
-def recursive_check_replaceme(value):
+def check_replaceme(value):
-    if isinstance(value, list):
-        return cv.Schema([recursive_check_replaceme])(value)
-    if isinstance(value, dict):
-        return cv.Schema({cv.valid: recursive_check_replaceme})(value)
-    if isinstance(value, ESPLiteralValue):
-        pass
     if isinstance(value, str) and value == "REPLACEME":
         raise cv.Invalid(
             "Found 'REPLACEME' in configuration, this is most likely an error. "
@@ -339,7 +333,86 @@ def recursive_check_replaceme(value):
             "If you want to use the literal REPLACEME string, "
             'please use "!literal REPLACEME"'
         )
-    return value
+
+
+def _build_list_index(lst):
+    index = OrderedDict()
+    extensions, removals = [], set()
+    for item in lst:
+        if item is None:
+            removals.add(None)
+            continue
+        item_id = None
+        if isinstance(item, dict) and (item_id := item.get(CONF_ID)):
+            if isinstance(item_id, Extend):
+                extensions.append(item)
+                continue
+            if isinstance(item_id, Remove):
+                removals.add(item_id.value)
+                continue
+        if not item_id or item_id in index:
+            # no id or duplicate -> pass through with identity-based key
+            item_id = id(item)
+        index[item_id] = item
+    return index, extensions, removals
+
+
+def resolve_extend_remove(value, is_key=None):
+    if isinstance(value, ESPLiteralValue):
+        return  # do not check inside literal blocks
+    if isinstance(value, list):
+        index, extensions, removals = _build_list_index(value)
+        if extensions or removals:
+            # Rebuild the original list after
+            # processing all extensions and removals
+            for item in extensions:
+                item_id = item[CONF_ID].value
+                if item_id in removals:
+                    continue
+                old = index.get(item_id)
+                if old is None:
+                    # Failed to find source for extension
+                    # Find index of item to show error at correct position
+                    i = next(
+                        (
+                            i
+                            for i, d in enumerate(value)
+                            if d.get(CONF_ID) == item[CONF_ID]
+                        )
+                    )
+                    with cv.prepend_path(i):
+                        raise cv.Invalid(
+                            f"Source for extension of ID '{item_id}' was not found."
+                        )
+                item[CONF_ID] = item_id
+                index[item_id] = merge_config(old, item)
+            for item_id in removals:
+                index.pop(item_id, None)
+
+            value[:] = index.values()
+
+        for i, item in enumerate(value):
+            with cv.prepend_path(i):
+                resolve_extend_remove(item, False)
+        return
+    if isinstance(value, dict):
+        removals = []
+        for k, v in value.items():
+            with cv.prepend_path(k):
+                if isinstance(v, Remove):
+                    removals.append(k)
+                    continue
+                resolve_extend_remove(k, True)
+                resolve_extend_remove(v, False)
+        for k in removals:
+            value.pop(k, None)
+        return
+    if is_key:
+        return  # do not check keys (yet)
+
+    check_replaceme(value)
+
+    return
 
 
 class ConfigValidationStep(abc.ABC):
@@ -437,19 +510,6 @@ class LoadValidationStep(ConfigValidationStep):
                 continue
             p_name = p_config.get("platform")
             if p_name is None:
-                p_id = p_config.get(CONF_ID)
-                if isinstance(p_id, Extend):
-                    result.add_str_error(
-                        f"Source for extension of ID '{p_id.value}' was not found.",
-                        path + [CONF_ID],
-                    )
-                    continue
-                if isinstance(p_id, Remove):
-                    result.add_str_error(
-                        f"Source for removal of ID '{p_id.value}' was not found.",
-                        path + [CONF_ID],
-                    )
-                    continue
                 result.add_str_error(
                     f"'{self.domain}' requires a 'platform' key but it was not specified.",
                     path,
@@ -934,9 +994,10 @@ def validate_config(
 
     CORE.raw_config = config
 
-    # 1.1. Check for REPLACEME special value
+    # 1.1. Resolve !extend and !remove and check for REPLACEME
+    # After this step, there will not be any Extend or Remove values in the config anymore
     try:
-        recursive_check_replaceme(config)
+        resolve_extend_remove(config)
     except vol.Invalid as err:
         result.add_error(err)
 

esphome/config_helpers.py

@@ -1,7 +1,6 @@
 from collections.abc import Callable
 
 from esphome.const import (
-    CONF_ID,
     CONF_LEVEL,
     CONF_LOGGER,
     KEY_CORE,
@@ -75,73 +74,28 @@ class Remove:
         return isinstance(b, Remove) and self.value == b.value
 
 
-def merge_config(full_old, full_new):
+def merge_config(old, new):
-    def merge(old, new):
+    if isinstance(new, Remove):
-        if isinstance(new, dict):
-            if not isinstance(old, dict):
-                return new
-            # Preserve OrderedDict type by copying to OrderedDict if either input is OrderedDict
-            if isinstance(old, OrderedDict) or isinstance(new, OrderedDict):
-                res = OrderedDict(old)
-            else:
-                res = old.copy()
-            for k, v in new.items():
-                if isinstance(v, Remove) and k in old:
-                    del res[k]
-                else:
-                    res[k] = merge(old[k], v) if k in old else v
-            return res
-        if isinstance(new, list):
-            if not isinstance(old, list):
-                return new
-            res = old.copy()
-            ids = {
-                v_id: i
-                for i, v in enumerate(res)
-                if isinstance(v, dict)
-                and (v_id := v.get(CONF_ID))
-                and isinstance(v_id, str)
-            }
-            extend_ids = {
-                v_id.value: i
-                for i, v in enumerate(res)
-                if isinstance(v, dict)
-                and (v_id := v.get(CONF_ID))
-                and isinstance(v_id, Extend)
-            }
-
-            ids_to_delete = []
-            for v in new:
-                if isinstance(v, dict) and (new_id := v.get(CONF_ID)):
-                    if isinstance(new_id, Extend):
-                        new_id = new_id.value
-                        if new_id in ids:
-                            v[CONF_ID] = new_id
-                            res[ids[new_id]] = merge(res[ids[new_id]], v)
-                            continue
-                    elif isinstance(new_id, Remove):
-                        new_id = new_id.value
-                        if new_id in ids:
-                            ids_to_delete.append(ids[new_id])
-                            continue
-                    elif (
-                        new_id in extend_ids
-                    ):  # When a package is extending a non-packaged item
-                        extend_res = res[extend_ids[new_id]]
-                        extend_res[CONF_ID] = new_id
-                        new_v = merge(v, extend_res)
-                        res[extend_ids[new_id]] = new_v
-                        continue
-                    else:
-                        ids[new_id] = len(res)
-                res.append(v)
-            return [v for i, v in enumerate(res) if i not in ids_to_delete]
-        if new is None:
-            return old
-
         return new
+    if isinstance(new, dict):
+        if not isinstance(old, dict):
+            return new
+        # Preserve OrderedDict type by copying to OrderedDict if either input is OrderedDict
+        if isinstance(old, OrderedDict) or isinstance(new, OrderedDict):
+            res = OrderedDict(old)
+        else:
+            res = old.copy()
+        for k, v in new.items():
+            res[k] = merge_config(old.get(k), v)
+        return res
+    if isinstance(new, list):
+        if not isinstance(old, list):
+            return new
+        return old + new
+    if new is None:
+        return old
 
-    return merge(full_old, full_new)
+    return new
 
 
 def filter_source_files_from_platform(

esphome/config_validation.py

@@ -24,7 +24,6 @@ import voluptuous as vol
 
 from esphome import core
 import esphome.codegen as cg
-from esphome.config_helpers import Extend, Remove
 from esphome.const import (
     ALLOWED_NAME_CHARS,
     CONF_AVAILABILITY,
@@ -624,12 +623,6 @@ def declare_id(type):
         if value is None:
             return core.ID(None, is_declaration=True, type=type)
 
-        if isinstance(value, Extend):
-            raise Invalid(f"Source for extension of ID '{value.value}' was not found.")
-
-        if isinstance(value, Remove):
-            raise Invalid(f"Source for Removal of ID '{value.value}' was not found.")
-
         return core.ID(validate_id_name(value), is_declaration=True, type=type)
 
     return validator

tests/component_tests/packages/test_packages.py

@@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 
 from esphome.components.packages import do_packages_pass
+from esphome.config import resolve_extend_remove
 from esphome.config_helpers import Extend, Remove
 import esphome.config_validation as cv
 from esphome.const import (
@@ -64,13 +65,20 @@ def fixture_basic_esphome():
     return {CONF_NAME: TEST_DEVICE_NAME, CONF_PLATFORM: TEST_PLATFORM}
 
 
+def packages_pass(config):
+    """Wrapper around packages_pass that also resolves Extend and Remove."""
+    config = do_packages_pass(config)
+    resolve_extend_remove(config)
+    return config
+
+
 def test_package_unused(basic_esphome, basic_wifi):
     """
     Ensures do_package_pass does not change a config if packages aren't used.
     """
     config = {CONF_ESPHOME: basic_esphome, CONF_WIFI: basic_wifi}
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == config
 
 
@@ -83,7 +91,7 @@ def test_package_invalid_dict(basic_esphome, basic_wifi):
     config = {CONF_ESPHOME: basic_esphome, CONF_PACKAGES: basic_wifi | {CONF_URL: ""}}
 
     with pytest.raises(cv.Invalid):
-        do_packages_pass(config)
+        packages_pass(config)
 
 
 def test_package_include(basic_wifi, basic_esphome):
@@ -99,7 +107,7 @@ def test_package_include(basic_wifi, basic_esphome):
 
     expected = {CONF_ESPHOME: basic_esphome, CONF_WIFI: basic_wifi}
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -124,7 +132,7 @@ def test_package_append(basic_wifi, basic_esphome):
         },
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -148,7 +156,7 @@ def test_package_override(basic_wifi, basic_esphome):
         },
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -177,7 +185,7 @@ def test_multiple_package_order():
         },
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -233,7 +241,7 @@ def test_package_list_merge():
         ]
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -311,7 +319,7 @@ def test_package_list_merge_by_id():
         ]
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -350,13 +358,13 @@ def test_package_merge_by_id_with_list():
         ]
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
 def test_package_merge_by_missing_id():
     """
-    Ensures that components with missing IDs are not merged.
+    Ensures that a validation error is thrown when trying to extend a missing ID.
     """
 
     config = {
@@ -379,25 +387,15 @@ def test_package_merge_by_missing_id():
         ],
     }
 
-    expected = {
+    error_raised = False
-        CONF_SENSOR: [
+    try:
-            {
+        packages_pass(config)
-                CONF_ID: TEST_SENSOR_ID_1,
+        assert False, "Expected validation error for missing ID"
-                CONF_FILTERS: [{CONF_MULTIPLY: 42.0}],
+    except cv.Invalid as err:
-            },
+        error_raised = True
-            {
+        assert err.path == [CONF_SENSOR, 2]
-                CONF_ID: TEST_SENSOR_ID_1,
-                CONF_FILTERS: [{CONF_MULTIPLY: 10.0}],
-            },
-            {
-                CONF_ID: Extend(TEST_SENSOR_ID_2),
-                CONF_FILTERS: [{CONF_OFFSET: 146.0}],
-            },
-        ]
-    }
 
-    actual = do_packages_pass(config)
+    assert error_raised
-    assert actual == expected
 
 
 def test_package_list_remove_by_id():
@@ -447,7 +445,7 @@ def test_package_list_remove_by_id():
         ]
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -493,7 +491,7 @@ def test_multiple_package_list_remove_by_id():
         ]
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -514,7 +512,7 @@ def test_package_dict_remove_by_id(basic_wifi, basic_esphome):
         CONF_ESPHOME: basic_esphome,
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -545,7 +543,6 @@ def test_package_remove_by_missing_id():
     }
 
     expected = {
-        "missing_key": Remove(),
         CONF_SENSOR: [
             {
                 CONF_ID: TEST_SENSOR_ID_1,
@@ -555,14 +552,10 @@ def test_package_remove_by_missing_id():
                 CONF_ID: TEST_SENSOR_ID_1,
                 CONF_FILTERS: [{CONF_MULTIPLY: 10.0}],
             },
-            {
-                CONF_ID: Remove(TEST_SENSOR_ID_2),
-                CONF_FILTERS: [{CONF_OFFSET: 146.0}],
-            },
         ],
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -634,7 +627,7 @@ def test_remote_packages_with_files_list(
         ]
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected
 
 
@@ -730,5 +723,5 @@ def test_remote_packages_with_files_and_vars(
         ]
     }
 
-    actual = do_packages_pass(config)
+    actual = packages_pass(config)
     assert actual == expected

tests/unit_tests/fixtures/substitutions/05-extend-remove.approved.yaml

@@ -0,0 +1,9 @@
+substitutions:
+  A: component1
+  B: component2
+  C: component3
+some_component:
+  - id: component1
+    value: 2
+  - id: component2
+    value: 5

tests/unit_tests/fixtures/substitutions/05-extend-remove.input.yaml

@@ -0,0 +1,22 @@
+substitutions:
+  A: component1
+  B: component2
+  C: component3
+
+packages:
+  - some_component:
+      - id: component1
+        value: 1
+      - id: !extend ${B}
+        value: 4
+      - id: !extend ${B}
+        value: 5
+      - id: component3
+        value: 6
+
+some_component:
+  - id: !extend ${A}
+    value: 2
+  - id: component2
+    value: 3
+  - id: !remove ${C}

tests/unit_tests/test_substitutions.py

@@ -4,6 +4,7 @@ from pathlib import Path
 
 from esphome import config as config_module, yaml_util
 from esphome.components import substitutions
+from esphome.config import resolve_extend_remove
 from esphome.config_helpers import merge_config
 from esphome.const import CONF_PACKAGES, CONF_SUBSTITUTIONS
 from esphome.core import CORE
@@ -81,6 +82,8 @@ def test_substitutions_fixtures(fixture_path):
 
             substitutions.do_substitution_pass(config, None)
 
+            resolve_extend_remove(config)
+
             # Also load expected using ESPHome's loader, or use {} if missing and DEV_MODE
             if expected_path.is_file():
                 expected = yaml_util.load_yaml(expected_path)