Make entity registry disabled_by an enum

homeassistant/helpers/entity_platform.py

@@ -42,7 +42,7 @@
     service,
 )
 from .device_registry import DeviceRegistry
-from .entity_registry import DISABLED_DEVICE, DISABLED_INTEGRATION, EntityRegistry
+from .entity_registry import EntityRegistry, RegistryEntryDisabler
 from .event import async_call_later, async_track_time_interval
 from .typing import ConfigType, DiscoveryInfoType
 
@@ -502,9 +502,9 @@ async def _async_add_entity(  # noqa: C901
                 except RequiredParameterMissing:
                     pass
 
-            disabled_by: str | None = None
+            disabled_by: RegistryEntryDisabler | None = None
             if not entity.entity_registry_enabled_default:
-                disabled_by = DISABLED_INTEGRATION
+                disabled_by = RegistryEntryDisabler.INTEGRATION
 
             entry = entity_registry.async_get_or_create(
                 self.domain,
@@ -526,7 +526,7 @@ async def _async_add_entity(  # noqa: C901
 
             if device and device.disabled and not entry.disabled:
                 entry = entity_registry.async_update_entity(
-                    entry.entity_id, disabled_by=DISABLED_DEVICE
+                    entry.entity_id, disabled_by=RegistryEntryDisabler.DEVICE
                 )
 
             entity.registry_entry = entry

homeassistant/helpers/entity_registry.py

@@ -39,8 +39,10 @@
 from homeassistant.exceptions import MaxLengthExceeded
 from homeassistant.helpers import device_registry as dr, storage
 from homeassistant.helpers.device_registry import EVENT_DEVICE_REGISTRY_UPDATED
+from homeassistant.helpers.frame import report
 from homeassistant.loader import bind_hass
 from homeassistant.util import slugify, uuid as uuid_util
+from homeassistant.util.enum import StrEnum
 from homeassistant.util.yaml import load_yaml
 
 from .typing import UNDEFINED, UndefinedType
@@ -53,11 +55,6 @@
 EVENT_ENTITY_REGISTRY_UPDATED = "entity_registry_updated"
 SAVE_DELAY = 10
 _LOGGER = logging.getLogger(__name__)
-DISABLED_CONFIG_ENTRY = "config_entry"
-DISABLED_DEVICE = "device"
-DISABLED_HASS = "hass"
-DISABLED_INTEGRATION = "integration"
-DISABLED_USER = "user"
 
 STORAGE_VERSION_MAJOR = 1
 STORAGE_VERSION_MINOR = 4
@@ -76,6 +73,24 @@
 }
 
 
+class RegistryEntryDisabler(StrEnum):
+    """What disabled a registry entry."""
+
+    CONFIG_ENTRY = "config_entry"
+    DEVICE = "device"
+    HASS = "hass"
+    INTEGRATION = "integration"
+    USER = "user"
+
+
+# DISABLED_* are deprecated, to be removed in 2022.3
+DISABLED_CONFIG_ENTRY = RegistryEntryDisabler.CONFIG_ENTRY.value
+DISABLED_DEVICE = RegistryEntryDisabler.DEVICE.value
+DISABLED_HASS = RegistryEntryDisabler.HASS.value
+DISABLED_INTEGRATION = RegistryEntryDisabler.INTEGRATION.value
+DISABLED_USER = RegistryEntryDisabler.USER.value
+
+
 @attr.s(slots=True, frozen=True)
 class RegistryEntry:
     """Entity Registry Entry."""
@@ -89,19 +104,7 @@ class RegistryEntry:
     device_class: str | None = attr.ib(default=None)
     device_id: str | None = attr.ib(default=None)
     domain: str = attr.ib(init=False, repr=False)
-    disabled_by: str | None = attr.ib(
-        default=None,
-        validator=attr.validators.in_(
-            (
-                DISABLED_CONFIG_ENTRY,
-                DISABLED_DEVICE,
-                DISABLED_HASS,
-                DISABLED_INTEGRATION,
-                DISABLED_USER,
-                None,
-            )
-        ),
-    )
+    disabled_by: RegistryEntryDisabler | None = attr.ib(default=None)
     entity_category: str | None = attr.ib(default=None)
     icon: str | None = attr.ib(default=None)
     id: str = attr.ib(factory=uuid_util.random_uuid_hex)
@@ -311,7 +314,7 @@ def async_get_or_create(
         known_object_ids: Iterable[str] | None = None,
         suggested_object_id: str | None = None,
         # To disable an entity if it gets created
-        disabled_by: str | None = None,
+        disabled_by: RegistryEntryDisabler | None = None,
         # Data that we want entry to have
         area_id: str | None = None,
         capabilities: Mapping[str, Any] | None = None,
@@ -357,12 +360,22 @@ def async_get_or_create(
             domain, suggested_object_id or f"{platform}_{unique_id}", known_object_ids
         )
 
-        if (
+        if isinstance(disabled_by, str) and not isinstance(
+            disabled_by, RegistryEntryDisabler
+        ):
+            report(  # type: ignore[unreachable]
+                "uses str for entity registry disabled_by. This is deprecated and will "
+                "stop working in Home Assistant 2022.3, it should be updated to use "
+                "RegistryEntryDisabler instead",
+                error_if_core=False,
+            )
+            disabled_by = RegistryEntryDisabler(disabled_by)
+        elif (
             disabled_by is None
             and config_entry
             and config_entry.pref_disable_new_entities
         ):
-            disabled_by = DISABLED_INTEGRATION
+            disabled_by = RegistryEntryDisabler.INTEGRATION
 
         entry = RegistryEntry(
             area_id=area_id,
@@ -429,7 +442,7 @@ def async_device_modified(self, event: Event) -> None:
                 self, event.data["device_id"], include_disabled_entities=True
             )
             for entity in entities:
-                if entity.disabled_by != DISABLED_DEVICE:
+                if entity.disabled_by is not RegistryEntryDisabler.DEVICE:
                     continue
                 self.async_update_entity(entity.entity_id, disabled_by=None)
             return
@@ -441,7 +454,9 @@ def async_device_modified(self, event: Event) -> None:
         # Fetch entities which are not already disabled
         entities = async_entries_for_device(self, event.data["device_id"])
         for entity in entities:
-            self.async_update_entity(entity.entity_id, disabled_by=DISABLED_DEVICE)
+            self.async_update_entity(
+                entity.entity_id, disabled_by=RegistryEntryDisabler.DEVICE
+            )
 
     @callback
     def async_update_entity(
@@ -451,7 +466,7 @@ def async_update_entity(
         area_id: str | None | UndefinedType = UNDEFINED,
         config_entry_id: str | None | UndefinedType = UNDEFINED,
         device_class: str | None | UndefinedType = UNDEFINED,
-        disabled_by: str | None | UndefinedType = UNDEFINED,
+        disabled_by: RegistryEntryDisabler | None | UndefinedType = UNDEFINED,
         entity_category: str | None | UndefinedType = UNDEFINED,
         icon: str | None | UndefinedType = UNDEFINED,
         name: str | None | UndefinedType = UNDEFINED,
@@ -490,7 +505,7 @@ def _async_update_entity(
         config_entry_id: str | None | UndefinedType = UNDEFINED,
         device_class: str | None | UndefinedType = UNDEFINED,
         device_id: str | None | UndefinedType = UNDEFINED,
-        disabled_by: str | None | UndefinedType = UNDEFINED,
+        disabled_by: RegistryEntryDisabler | None | UndefinedType = UNDEFINED,
         entity_category: str | None | UndefinedType = UNDEFINED,
         icon: str | None | UndefinedType = UNDEFINED,
         name: str | None | UndefinedType = UNDEFINED,
@@ -508,6 +523,17 @@ def _async_update_entity(
         new_values = {}  # Dict with new key/value pairs
         old_values = {}  # Dict with old key/value pairs
 
+        if isinstance(disabled_by, str) and not isinstance(
+            disabled_by, RegistryEntryDisabler
+        ):
+            report(  # type: ignore[unreachable]
+                "uses str for entity registry disabled_by. This is deprecated and will "
+                "stop working in Home Assistant 2022.3, it should be updated to use "
+                "RegistryEntryDisabler instead",
+                error_if_core=False,
+            )
+            disabled_by = RegistryEntryDisabler(disabled_by)
+
         for attr_name, value in (
             ("area_id", area_id),
             ("capabilities", capabilities),
@@ -597,7 +623,9 @@ async def async_load(self) -> None:
                     config_entry_id=entity["config_entry_id"],
                     device_class=entity["device_class"],
                     device_id=entity["device_id"],
-                    disabled_by=entity["disabled_by"],
+                    disabled_by=RegistryEntryDisabler(entity["disabled_by"])
+                    if entity["disabled_by"]
+                    else None,
                     entity_category=entity["entity_category"],
                     entity_id=entity["entity_id"],
                     icon=entity["icon"],
@@ -739,7 +767,7 @@ def async_config_entry_disabled_by_changed(
 
     if not config_entry.disabled_by:
         for entity in entities:
-            if entity.disabled_by != DISABLED_CONFIG_ENTRY:
+            if entity.disabled_by is not RegistryEntryDisabler.CONFIG_ENTRY:
                 continue
             registry.async_update_entity(entity.entity_id, disabled_by=None)
         return
@@ -749,7 +777,7 @@ def async_config_entry_disabled_by_changed(
             # Entity already disabled, do not overwrite
             continue
         registry.async_update_entity(
-            entity.entity_id, disabled_by=DISABLED_CONFIG_ENTRY
+            entity.entity_id, disabled_by=RegistryEntryDisabler.CONFIG_ENTRY
         )
 
 

tests/helpers/test_entity_registry.py

@@ -1081,3 +1081,16 @@ def test_entity_registry_items():
     assert entities.get_entry(entry1.id) is None
     assert entities.get_entity_id(("test", "hue", "2345")) is None
     assert entities.get_entry(entry2.id) is None
+
+
+async def test_deprecated_disabled_by_str(hass, registry, caplog):
+    """Test deprecated str use of disabled_by converts to enum and logs a warning."""
+    entry = registry.async_get_or_create(
+        "light",
+        "hue",
+        "5678",
+        disabled_by=er.RegistryEntryDisabler.USER.value,
+    )
+
+    assert entry.disabled_by is er.RegistryEntryDisabler.USER
+    assert " str for entity registry disabled_by. This is deprecated " in caplog.text