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>
This commit is contained in:
epenet 2022-01-24 12:51:11 +01:00 committed by GitHub
parent b03ae66254
commit 32d5c8e8ac
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 44 additions and 8 deletions

View File

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

View File

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

View File

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

View File

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