Migrate entity categories to StrEnum (#60720)

homeassistant/const.py

@@ -703,16 +703,11 @@ PRECISION_TENTHS: Final = 0.1
 # cloud, alexa, or google_home components
 CLOUD_NEVER_EXPOSED_ENTITIES: Final[list[str]] = ["group.all_locks"]
 
-# Config: An entity which allows changing the configuration of a device
+# ENTITY_CATEGOR* below are deprecated as of 2021.12
+# use the EntityCategory enum instead.
 ENTITY_CATEGORY_CONFIG: Final = "config"
-# Diagnostic: An entity exposing some configuration parameter or diagnostics of a device
 ENTITY_CATEGORY_DIAGNOSTIC: Final = "diagnostic"
-# System: An entity which is not useful for the user to interact with
 ENTITY_CATEGORY_SYSTEM: Final = "system"
-
-# Entity categories which will:
-# - Not be exposed to cloud, alexa, or google_home components
-# - Not be included in indirect service calls to devices or areas
 ENTITY_CATEGORIES: Final[list[str]] = [
     ENTITY_CATEGORY_CONFIG,
     ENTITY_CATEGORY_DIAGNOSTIC,

homeassistant/helpers/entity.py

@@ -43,6 +43,7 @@ from homeassistant.helpers.event import Event, async_track_entity_registry_updat
 from homeassistant.helpers.typing import StateType
 from homeassistant.loader import bind_hass
 from homeassistant.util import dt as dt_util, ensure_unique_string, slugify
+from homeassistant.util.enum import StrEnum
 
 _LOGGER = logging.getLogger(__name__)
 SLOW_UPDATE_WARNING = 10
@@ -179,6 +180,24 @@ class DeviceInfo(TypedDict, total=False):
     via_device: tuple[str, str]
 
 
+class EntityCategory(StrEnum):
+    """Category of an entity.
+
+    An entity with a category will:
+    - Not be exposed to cloud, Alexa, or Google Assistant components
+    - Not be included in indirect service calls to devices or areas
+    """
+
+    # Config: An entity which allows changing the configuration of a device
+    CONFIG = "config"
+
+    # Diagnostic: An entity exposing some configuration parameter or diagnostics of a device
+    DIAGNOSTIC = "diagnostic"
+
+    # System: An entity which is not useful for the user to interact with
+    SYSTEM = "system"
+
+
 @dataclass
 class EntityDescription:
     """A class that describes Home Assistant entities."""
@@ -187,7 +206,9 @@ class EntityDescription:
     key: str
 
     device_class: str | None = None
-    entity_category: Literal["config", "diagnostic", "system"] | None = None
+    entity_category: EntityCategory | Literal[
+        "config", "diagnostic", "system"
+    ] | None = None
     entity_registry_enabled_default: bool = True
     force_update: bool = False
     icon: str | None = None
@@ -246,7 +267,7 @@ class Entity(ABC):
     _attr_context_recent_time: timedelta = timedelta(seconds=5)
     _attr_device_class: str | None
     _attr_device_info: DeviceInfo | None = None
-    _attr_entity_category: str | None
+    _attr_entity_category: EntityCategory | str | None
     _attr_entity_picture: str | None = None
     _attr_entity_registry_enabled_default: bool
     _attr_extra_state_attributes: MutableMapping[str, Any]
@@ -414,7 +435,7 @@ class Entity(ABC):
         return self._attr_attribution
 
     @property
-    def entity_category(self) -> str | None:
+    def entity_category(self) -> EntityCategory | str | None:
         """Return the category of the entity, if any."""
         if hasattr(self, "_attr_entity_category"):
             return self._attr_entity_category

tests/helpers/test_entity.py

@@ -819,13 +819,13 @@ async def test_entity_category_property(hass):
         key="abc", entity_category="ignore_me"
     )
     mock_entity1.entity_id = "hello.world"
-    mock_entity1._attr_entity_category = "config"
+    mock_entity1._attr_entity_category = entity.EntityCategory.CONFIG
     assert mock_entity1.entity_category == "config"
 
     mock_entity2 = entity.Entity()
     mock_entity2.hass = hass
     mock_entity2.entity_description = entity.EntityDescription(
-        key="abc", entity_category="config"
+        key="abc", entity_category=entity.EntityCategory.CONFIG
     )
     mock_entity2.entity_id = "hello.world"
     assert mock_entity2.entity_category == "config"