[psram] Make schema extractable with per-variant options (#16949)

Co-authored-by: J. Nick Koston <nick@koston.org>

esphome/components/esp32/__init__.py

@@ -1,3 +1,4 @@
+from collections.abc import Callable, Iterable
 import contextlib
 from dataclasses import dataclass
 import itertools
@@ -6,6 +7,7 @@ import os
 from pathlib import Path
 import re
 import subprocess
+from typing import Any
 
 from esphome import yaml_util
 import esphome.codegen as cg
@@ -52,6 +54,7 @@ from esphome.coroutine import CoroPriority, coroutine_with_priority
 from esphome.espidf.component import generate_idf_components
 import esphome.final_validate as fv
 from esphome.helpers import copy_file_if_changed, rmtree, write_file_if_changed
+from esphome.schema_extractors import SCHEMA_EXTRACT, schema_extractor
 from esphome.types import ConfigType
 from esphome.writer import clean_build, clean_cmake_cache
 
@@ -496,6 +499,32 @@ def get_esp32_variant(core_obj=None):
     return (core_obj or CORE).data[KEY_ESP32][KEY_VARIANT]
 
 
+def variant_filtered_enum(
+    by_variant: dict[str, Iterable[Any]], **kwargs: Any
+) -> Callable[[Any], Any]:
+    """Build a ``one_of`` validator whose valid set depends on the active variant.
+
+    ``by_variant`` maps each ESP32 variant constant to the iterable of values that
+    are valid on that variant. At validation time the value is checked against the
+    set allowed for the current target variant. For schema extraction the inverted
+    ``{value: [variants, ...]}`` map is returned instead, so the language-schema
+    dump can tag every option with the variants that accept it and frontends can
+    filter to the user's selected variant.
+    """
+    by_value: dict[str, list[str]] = {}
+    for variant, values in by_variant.items():
+        for value in values:
+            by_value.setdefault(str(value), []).append(variant)
+
+    @schema_extractor("variant_enum")
+    def validator(value: Any) -> Any:
+        if value is SCHEMA_EXTRACT:
+            return by_value
+        return cv.one_of(*by_variant.get(get_esp32_variant(), ()), **kwargs)(value)
+
+    return validator
+
+
 def get_board(core_obj=None):
     return (core_obj or CORE).data[KEY_ESP32][KEY_BOARD]
 

esphome/components/psram/__init__.py

@@ -16,6 +16,7 @@ from esphome.components.esp32 import (
     add_idf_sdkconfig_option,
     get_esp32_variant,
     idf_version,
+    variant_filtered_enum,
 )
 import esphome.config_validation as cv
 from esphome.const import (
@@ -29,6 +30,7 @@ from esphome.const import (
 )
 from esphome.core import CORE
 import esphome.final_validate as fv
+from esphome.types import ConfigType
 
 CODEOWNERS = ["@esphome/core"]
 DOMAIN = "psram"
@@ -70,6 +72,11 @@ SPIRAM_SPEEDS = {
     VARIANT_ESP32P4: (20, 100, 200),
 }
 
+SPIRAM_SPEEDS_MHZ = {
+    variant: tuple(f"{speed}MHZ" for speed in speeds)
+    for variant, speeds in SPIRAM_SPEEDS.items()
+}
+
 
 def supported() -> bool:
     if not CORE.is_esp32:
@@ -145,15 +152,23 @@ def validate_psram_mode(config):
     return config
 
 
-def get_config_schema(config):
+def _set_variant_defaults(config: ConfigType) -> ConfigType:
+    """Resolve variant-dependent defaults before the static schema validates.
+
+    The set of valid ``mode``/``speed`` values is variant-specific (enforced by
+    ``variant_filtered_enum`` in the schema below); this only supplies the default
+    when the user omits the option. ``mode`` has no single default on chips that
+    support more than one mode, so selection is required there.
+    """
     variant = get_esp32_variant()
-    speeds = [f"{s}MHZ" for s in SPIRAM_SPEEDS.get(variant, [])]
-    if not speeds:
+    modes = SPIRAM_MODES.get(variant)
+    speeds = SPIRAM_SPEEDS.get(variant)
+    if not modes or not speeds:
         raise cv.Invalid("PSRAM is not supported on this chip")
-    modes = SPIRAM_MODES[variant]
-    if CONF_MODE not in config and len(modes) != 1:
-        raise (
-            cv.Invalid(
+    config = config.copy()
+    if CONF_MODE not in config:
+        if len(modes) != 1:
+            raise cv.Invalid(
                 textwrap.dedent(
                     f"""
                         {variant} requires PSRAM mode selection; one of {", ".join(modes)}
@@ -161,20 +176,27 @@ def get_config_schema(config):
                     """
                 )
             )
-        )
-    return cv.Schema(
+        config[CONF_MODE] = modes[0]
+    if CONF_SPEED not in config:
+        config[CONF_SPEED] = f"{speeds[0]}MHZ"
+    return config
+
+
+CONFIG_SCHEMA = cv.All(
+    _set_variant_defaults,
+    cv.Schema(
         {
             cv.GenerateID(): cv.declare_id(PsramComponent),
-            cv.Optional(CONF_MODE, default=modes[0]): cv.one_of(*modes, lower=True),
+            cv.Optional(CONF_MODE): variant_filtered_enum(SPIRAM_MODES, lower=True),
             cv.Optional(CONF_ENABLE_ECC, default=False): cv.boolean,
-            cv.Optional(CONF_SPEED, default=speeds[0]): cv.one_of(*speeds, upper=True),
+            cv.Optional(CONF_SPEED): variant_filtered_enum(
+                SPIRAM_SPEEDS_MHZ, upper=True
+            ),
             cv.Optional(CONF_DISABLED, default=False): cv.boolean,
             cv.Optional(CONF_IGNORE_NOT_FOUND, default=True): cv.boolean,
         }
-    )(config)
-
-
-CONFIG_SCHEMA = get_config_schema
+    ),
+)
 
 
 def _store_psram_guaranteed(config):