From 8fd60dbd51a5a192e4ed1e2f250d7c618a82b1aa Mon Sep 17 00:00:00 2001
From: Dave T <17680170+davet2001@users.noreply.github.com>
Date: Mon, 27 Dec 2021 17:23:08 +0000
Subject: [PATCH] Refactor entity_category str types (#62790)

---
 homeassistant/components/mqtt/mixins.py    |  3 ++-
 homeassistant/components/neato/sensor.py   |  2 +-
 homeassistant/components/shelly/entity.py  | 14 +++++++-------
 homeassistant/components/tasmota/sensor.py |  2 +-
 4 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/mqtt/mixins.py b/homeassistant/components/mqtt/mixins.py
index 713f0e5c030..9b6d4323433 100644
--- a/homeassistant/components/mqtt/mixins.py
+++ b/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)
 
diff --git a/homeassistant/components/neato/sensor.py b/homeassistant/components/neato/sensor.py
index 7cc2d0f171a..2b20158ff8b 100644
--- a/homeassistant/components/neato/sensor.py
+++ b/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
 
diff --git a/homeassistant/components/shelly/entity.py b/homeassistant/components/shelly/entity.py
index 65ce8eeba56..d4f26c7767d 100644
--- a/homeassistant/components/shelly/entity.py
+++ b/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
 
diff --git a/homeassistant/components/tasmota/sensor.py b/homeassistant/components/tasmota/sensor.py
index ece45b9dfd7..feb15cd5639 100644
--- a/homeassistant/components/tasmota/sensor.py
+++ b/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