Add type hint to RegistryEntry.entity_category (#64618)

Co-authored-by: Dave T <17680170+davet2001@users.noreply.github.com>
Co-authored-by: epenet <epenet@users.noreply.github.com>

homeassistant/components/goodwe/sensor.py

@@ -24,14 +24,13 @@ from homeassistant.const import (
     ELECTRIC_CURRENT_AMPERE,
     ELECTRIC_POTENTIAL_VOLT,
     ENERGY_KILO_WATT_HOUR,
-    ENTITY_CATEGORY_DIAGNOSTIC,
     FREQUENCY_HERTZ,
     PERCENTAGE,
     POWER_WATT,
     TEMP_CELSIUS,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.entity import DeviceInfo, EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
@@ -159,7 +158,7 @@ class InverterSensor(CoordinatorEntity, SensorEntity):
         self._attr_unique_id = f"{DOMAIN}-{sensor.id_}-{inverter.serial_number}"
         self._attr_device_info = device_info
         self._attr_entity_category = (
-            ENTITY_CATEGORY_DIAGNOSTIC if sensor.id_ not in _MAIN_SENSORS else None
+            EntityCategory.DIAGNOSTIC if sensor.id_ not in _MAIN_SENSORS else None
         )
         self.entity_description = _DESCRIPTIONS.get(sensor.unit, DIAG_SENSOR)
         if not self.entity_description.native_unit_of_measurement:

homeassistant/components/smartthings/sensor.py

@@ -628,7 +628,7 @@ class SmartThingsSensor(SmartThingsEntity, SensorEntity):
         default_unit: str,
         device_class: str,
         state_class: str | None,
-        entity_category: str | None,
+        entity_category: EntityCategory | None,
     ) -> None:
         """Init the class."""
         super().__init__(device)

homeassistant/helpers/entity.py

@@ -44,6 +44,7 @@ from . import entity_registry as er
 from .device_registry import DeviceEntryType
 from .entity_platform import EntityPlatform
 from .event import async_track_entity_registry_updated_event
+from .frame import report
 from .typing import StateType
 
 _LOGGER = logging.getLogger(__name__)
@@ -200,6 +201,24 @@ class EntityCategory(StrEnum):
     SYSTEM = "system"
 
 
+def convert_to_entity_category(
+    value: EntityCategory | str | None,
+) -> EntityCategory | None:
+    """Force incoming entity_category to be an enum."""
+
+    if value is None:
+        return value
+
+    if not isinstance(value, EntityCategory):
+        report(
+            "An entity_category should only be assigned an enum.  Strings or other assignments are deprecated. Value %s is type %s"
+            % (value, type(value)),
+            error_if_core=False,
+        )
+        return EntityCategory(value)
+    return value
+
+
 @dataclass
 class EntityDescription:
     """A class that describes Home Assistant entities."""
@@ -208,6 +227,7 @@ class EntityDescription:
     key: str
 
     device_class: str | None = None
+    # Type string is deprecated as of 2021.12, use EntityCategory
     entity_category: EntityCategory | Literal[
         "config", "diagnostic", "system"
     ] | None = None
@@ -272,7 +292,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: EntityCategory | str | None
+    _attr_entity_category: EntityCategory | None
     _attr_entity_picture: str | None = None
     _attr_entity_registry_enabled_default: bool
     _attr_extra_state_attributes: MutableMapping[str, Any]
@@ -439,6 +459,7 @@ class Entity(ABC):
         """Return the attribution."""
         return self._attr_attribution
 
+    # Type str is deprecated as of 2021.12, use EntityCategory
     @property
     def entity_category(self) -> EntityCategory | str | None:
         """Return the category of the entity, if any."""

homeassistant/helpers/entity_registry.py

@@ -50,6 +50,8 @@ from .typing import UNDEFINED, UndefinedType
 if TYPE_CHECKING:
     from homeassistant.config_entries import ConfigEntry
 
+    from .entity import EntityCategory
+
 PATH_REGISTRY = "entity_registry.yaml"
 DATA_REGISTRY = "entity_registry"
 EVENT_ENTITY_REGISTRY_UPDATED = "entity_registry_updated"
@@ -91,6 +93,16 @@ DISABLED_INTEGRATION = RegistryEntryDisabler.INTEGRATION.value
 DISABLED_USER = RegistryEntryDisabler.USER.value
 
 
+def _convert_to_entity_category(
+    value: EntityCategory | str | None,
+) -> EntityCategory | None:
+    """Force incoming entity_category to be an enum."""
+    # pylint: disable=import-outside-toplevel
+    from .entity import convert_to_entity_category
+
+    return convert_to_entity_category(value)
+
+
 @attr.s(slots=True, frozen=True)
 class RegistryEntry:
     """Entity Registry Entry."""
@@ -105,7 +117,9 @@ class RegistryEntry:
     device_id: str | None = attr.ib(default=None)
     domain: str = attr.ib(init=False, repr=False)
     disabled_by: RegistryEntryDisabler | None = attr.ib(default=None)
-    entity_category: str | None = attr.ib(default=None)
+    entity_category: EntityCategory | None = attr.ib(
+        default=None, converter=_convert_to_entity_category
+    )
     icon: str | None = attr.ib(default=None)
     id: str = attr.ib(factory=uuid_util.random_uuid_hex)
     name: str | None = attr.ib(default=None)
@@ -320,7 +334,8 @@ class EntityRegistry:
         capabilities: Mapping[str, Any] | None = None,
         config_entry: ConfigEntry | None = None,
         device_id: str | None = None,
-        entity_category: str | None = None,
+        # Type str (ENTITY_CATEG*) is deprecated as of 2021.12, use EntityCategory
+        entity_category: EntityCategory | str | None = None,
         original_device_class: str | None = None,
         original_icon: str | None = None,
         original_name: str | None = None,
@@ -469,7 +484,8 @@ class EntityRegistry:
         device_class: str | None | UndefinedType = UNDEFINED,
         device_id: str | None | UndefinedType = UNDEFINED,
         disabled_by: RegistryEntryDisabler | None | UndefinedType = UNDEFINED,
-        entity_category: str | None | UndefinedType = UNDEFINED,
+        # Type str (ENTITY_CATEG*) is deprecated as of 2021.12, use EntityCategory
+        entity_category: EntityCategory | str | None | UndefinedType = UNDEFINED,
         icon: str | None | UndefinedType = UNDEFINED,
         name: str | None | UndefinedType = UNDEFINED,
         new_entity_id: str | UndefinedType = UNDEFINED,