Refactor entity_category str types (#62790)

This commit is contained in:
Dave T 2021-12-27 17:23:08 +00:00 committed by GitHub
parent 377b0efc60
commit 8fd60dbd51
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 11 additions and 10 deletions

View File

@ -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)

View File

@ -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

View File

@ -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

View File

@ -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