Refactor entity_category str types (#62790)

homeassistant/components/mqtt/mixins.py

@@ -33,6 +33,7 @@ from homeassistant.helpers.entity import (
     ENTITY_CATEGORIES_SCHEMA,
     DeviceInfo,
     Entity,
+    EntityCategory,
     async_generate_entity_id,
 )
 from homeassistant.helpers.typing import ConfigType
@@ -695,7 +696,7 @@ class MqttEntity(
         return self._config[CONF_ENABLED_BY_DEFAULT]
 
     @property
-    def entity_category(self) -> str | None:
+    def entity_category(self) -> EntityCategory | str | None:
         """Return the entity category if any."""
         return self._config.get(CONF_ENTITY_CATEGORY)
 

homeassistant/components/neato/sensor.py

@@ -84,7 +84,7 @@ class NeatoSensor(SensorEntity):
         return SensorDeviceClass.BATTERY
 
     @property
-    def entity_category(self) -> str:
+    def entity_category(self) -> EntityCategory:
         """Device entity category."""
         return EntityCategory.DIAGNOSTIC
 

homeassistant/components/shelly/entity.py

@@ -19,7 +19,7 @@ from homeassistant.helpers import (
     entity_registry,
     update_coordinator,
 )
-from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.entity import DeviceInfo, EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import StateType
@@ -237,7 +237,7 @@ class BlockAttributeDescription:
     # Callable (settings, block), return true if entity should be removed
     removal_condition: Callable[[dict, Block], bool] | None = None
     extra_state_attributes: Callable[[Block], dict | None] | None = None
-    entity_category: str | None = None
+    entity_category: EntityCategory | None = None
 
 
 @dataclass
@@ -256,7 +256,7 @@ class RpcAttributeDescription:
     available: Callable[[dict], bool] | None = None
     removal_condition: Callable[[dict, str], bool] | None = None
     extra_state_attributes: Callable[[dict, dict], dict | None] | None = None
-    entity_category: str | None = None
+    entity_category: EntityCategory | None = None
 
 
 @dataclass
@@ -271,7 +271,7 @@ class RestAttributeDescription:
     state_class: str | None = None
     default_enabled: bool = True
     extra_state_attributes: Callable[[dict], dict | None] | None = None
-    entity_category: str | None = None
+    entity_category: EntityCategory | None = None
 
 
 class ShellyBlockEntity(entity.Entity):
@@ -471,7 +471,7 @@ class ShellyBlockAttributeEntity(ShellyBlockEntity, entity.Entity):
         return self.description.extra_state_attributes(self.block)
 
     @property
-    def entity_category(self) -> str | None:
+    def entity_category(self) -> EntityCategory | None:
         """Return category of entity."""
         return self.description.entity_category
 
@@ -548,7 +548,7 @@ class ShellyRestAttributeEntity(update_coordinator.CoordinatorEntity):
         return self.description.extra_state_attributes(self.wrapper.device.status)
 
     @property
-    def entity_category(self) -> str | None:
+    def entity_category(self) -> EntityCategory | None:
         """Return category of entity."""
         return self.description.entity_category
 
@@ -614,7 +614,7 @@ class ShellyRpcAttributeEntity(ShellyRpcEntity, entity.Entity):
         )
 
     @property
-    def entity_category(self) -> str | None:
+    def entity_category(self) -> EntityCategory | None:
         """Return category of entity."""
         return self.description.entity_category
 

homeassistant/components/tasmota/sensor.py

@@ -274,7 +274,7 @@ class TasmotaSensor(TasmotaAvailability, TasmotaDiscoveryUpdate, SensorEntity):
         return class_or_icon.get(STATE_CLASS)
 
     @property
-    def entity_category(self) -> str | None:
+    def entity_category(self) -> EntityCategory | None:
         """Return the category of the entity, if any."""
         if self._tasmota_entity.quantity in status_sensor.SENSORS:
             return EntityCategory.DIAGNOSTIC