From b84829f70fcd8dfb73cbd9829cf962fc029f0d8e Mon Sep 17 00:00:00 2001
From: Allen Porter <allen@thebends.org>
Date: Mon, 13 May 2024 21:07:39 -0700
Subject: [PATCH] Import and cache supported feature enum flags only when
 needed (#117270)

* Import and cache supported feature enum flags only when needed

* Add comment aboud being loaded from executor.

---------

Co-authored-by: J. Nick Koston <nick@koston.org>
---
 homeassistant/helpers/selector.py | 67 +++++++------------------------
 tests/helpers/test_selector.py    |  2 +
 2 files changed, 16 insertions(+), 53 deletions(-)

diff --git a/homeassistant/helpers/selector.py b/homeassistant/helpers/selector.py
index a45ba2d1129..01521556453 100644
--- a/homeassistant/helpers/selector.py
+++ b/homeassistant/helpers/selector.py
@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping, Sequence
-from enum import IntFlag, StrEnum
+from enum import StrEnum
 from functools import cache
+import importlib
 from typing import Any, Generic, Literal, Required, TypedDict, TypeVar, cast
 from uuid import UUID
 
@@ -82,63 +83,23 @@ class Selector(Generic[_T]):
 
 
 @cache
-def _entity_features() -> dict[str, type[IntFlag]]:
-    """Return a cached lookup of entity feature enums."""
-    # pylint: disable=import-outside-toplevel
-    from homeassistant.components.alarm_control_panel import (
-        AlarmControlPanelEntityFeature,
-    )
-    from homeassistant.components.calendar import CalendarEntityFeature
-    from homeassistant.components.camera import CameraEntityFeature
-    from homeassistant.components.climate import ClimateEntityFeature
-    from homeassistant.components.cover import CoverEntityFeature
-    from homeassistant.components.fan import FanEntityFeature
-    from homeassistant.components.humidifier import HumidifierEntityFeature
-    from homeassistant.components.lawn_mower import LawnMowerEntityFeature
-    from homeassistant.components.light import LightEntityFeature
-    from homeassistant.components.lock import LockEntityFeature
-    from homeassistant.components.media_player import MediaPlayerEntityFeature
-    from homeassistant.components.notify import NotifyEntityFeature
-    from homeassistant.components.remote import RemoteEntityFeature
-    from homeassistant.components.siren import SirenEntityFeature
-    from homeassistant.components.todo import TodoListEntityFeature
-    from homeassistant.components.update import UpdateEntityFeature
-    from homeassistant.components.vacuum import VacuumEntityFeature
-    from homeassistant.components.valve import ValveEntityFeature
-    from homeassistant.components.water_heater import WaterHeaterEntityFeature
-    from homeassistant.components.weather import WeatherEntityFeature
+def _entity_feature_flag(domain: str, enum_name: str, feature_name: str) -> int:
+    """Return a cached lookup of an entity feature enum.
 
-    return {
-        "AlarmControlPanelEntityFeature": AlarmControlPanelEntityFeature,
-        "CalendarEntityFeature": CalendarEntityFeature,
-        "CameraEntityFeature": CameraEntityFeature,
-        "ClimateEntityFeature": ClimateEntityFeature,
-        "CoverEntityFeature": CoverEntityFeature,
-        "FanEntityFeature": FanEntityFeature,
-        "HumidifierEntityFeature": HumidifierEntityFeature,
-        "LawnMowerEntityFeature": LawnMowerEntityFeature,
-        "LightEntityFeature": LightEntityFeature,
-        "LockEntityFeature": LockEntityFeature,
-        "MediaPlayerEntityFeature": MediaPlayerEntityFeature,
-        "NotifyEntityFeature": NotifyEntityFeature,
-        "RemoteEntityFeature": RemoteEntityFeature,
-        "SirenEntityFeature": SirenEntityFeature,
-        "TodoListEntityFeature": TodoListEntityFeature,
-        "UpdateEntityFeature": UpdateEntityFeature,
-        "VacuumEntityFeature": VacuumEntityFeature,
-        "ValveEntityFeature": ValveEntityFeature,
-        "WaterHeaterEntityFeature": WaterHeaterEntityFeature,
-        "WeatherEntityFeature": WeatherEntityFeature,
-    }
+    This will import a module from disk and is run from an executor when
+    loading the services schema files.
+    """
+    module = importlib.import_module(f"homeassistant.components.{domain}")
+    enum = getattr(module, enum_name)
+    feature = getattr(enum, feature_name)
+    return cast(int, feature.value)
 
 
 def _validate_supported_feature(supported_feature: str) -> int:
     """Validate a supported feature and resolve an enum string to its value."""
 
-    known_entity_features = _entity_features()
-
     try:
-        _, enum, feature = supported_feature.split(".", 2)
+        domain, enum, feature = supported_feature.split(".", 2)
     except ValueError as exc:
         raise vol.Invalid(
             f"Invalid supported feature '{supported_feature}', expected "
@@ -146,8 +107,8 @@ def _validate_supported_feature(supported_feature: str) -> int:
         ) from exc
 
     try:
-        return cast(int, getattr(known_entity_features[enum], feature).value)
-    except (AttributeError, KeyError) as exc:
+        return _entity_feature_flag(domain, enum, feature)
+    except (ModuleNotFoundError, AttributeError) as exc:
         raise vol.Invalid(f"Unknown supported feature '{supported_feature}'") from exc
 
 
diff --git a/tests/helpers/test_selector.py b/tests/helpers/test_selector.py
index 8864edc7386..5e6209f2c6c 100644
--- a/tests/helpers/test_selector.py
+++ b/tests/helpers/test_selector.py
@@ -282,6 +282,8 @@ def test_entity_selector_schema(schema, valid_selections, invalid_selections) ->
         {"filter": [{"supported_features": ["blah"]}]},
         # Unknown feature enum
         {"filter": [{"supported_features": ["blah.FooEntityFeature.blah"]}]},
+        # Unknown feature enum
+        {"filter": [{"supported_features": ["light.FooEntityFeature.blah"]}]},
         # Unknown feature enum member
         {"filter": [{"supported_features": ["light.LightEntityFeature.blah"]}]},
     ],