Use new enums in esphome (#61391)

Co-authored-by: epenet <epenet@users.noreply.github.com>
This commit is contained in:
epenet 2021-12-10 09:56:20 +01:00 committed by GitHub
parent 412e531096
commit e50c00ea06
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 21 deletions

View File

@ -13,7 +13,7 @@ from aioesphomeapi import (
APIIntEnum, APIIntEnum,
APIVersion, APIVersion,
DeviceInfo as EsphomeDeviceInfo, DeviceInfo as EsphomeDeviceInfo,
EntityCategory, EntityCategory as EsphomeEntityCategory,
EntityInfo, EntityInfo,
EntityState, EntityState,
HomeassistantServiceCall, HomeassistantServiceCall,
@ -33,8 +33,6 @@ from homeassistant.const import (
CONF_MODE, CONF_MODE,
CONF_PASSWORD, CONF_PASSWORD,
CONF_PORT, CONF_PORT,
ENTITY_CATEGORY_CONFIG,
ENTITY_CATEGORY_DIAGNOSTIC,
EVENT_HOMEASSISTANT_STOP, EVENT_HOMEASSISTANT_STOP,
) )
from homeassistant.core import Event, HomeAssistant, ServiceCall, State, callback from homeassistant.core import Event, HomeAssistant, ServiceCall, State, callback
@ -43,7 +41,7 @@ from homeassistant.helpers import template
import homeassistant.helpers.config_validation as cv import homeassistant.helpers.config_validation as cv
import homeassistant.helpers.device_registry as dr import homeassistant.helpers.device_registry as dr
from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import DeviceInfo, Entity from homeassistant.helpers.entity import DeviceInfo, Entity, EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import async_track_state_change_event from homeassistant.helpers.event import async_track_state_change_event
from homeassistant.helpers.json import JSONEncoder from homeassistant.helpers.json import JSONEncoder
@ -647,11 +645,13 @@ class EsphomeEnumMapper(Generic[_EnumT, _ValT]):
ICON_SCHEMA = vol.Schema(cv.icon) ICON_SCHEMA = vol.Schema(cv.icon)
ENTITY_CATEGORIES: EsphomeEnumMapper[EntityCategory, str | None] = EsphomeEnumMapper( ENTITY_CATEGORIES: EsphomeEnumMapper[
EsphomeEntityCategory, EntityCategory | None
] = EsphomeEnumMapper(
{ {
EntityCategory.NONE: None, EsphomeEntityCategory.NONE: None,
EntityCategory.CONFIG: ENTITY_CATEGORY_CONFIG, EsphomeEntityCategory.CONFIG: EntityCategory.CONFIG,
EntityCategory.DIAGNOSTIC: ENTITY_CATEGORY_DIAGNOSTIC, EsphomeEntityCategory.DIAGNOSTIC: EntityCategory.DIAGNOSTIC,
} }
) )
@ -799,7 +799,7 @@ class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
return not self._static_info.disabled_by_default return not self._static_info.disabled_by_default
@property @property
def entity_category(self) -> str | None: def entity_category(self) -> EntityCategory | None:
"""Return the category of the entity, if any.""" """Return the category of the entity, if any."""
if not self._static_info.entity_category: if not self._static_info.entity_category:
return None return None

View File

@ -7,18 +7,17 @@ import math
from aioesphomeapi import ( from aioesphomeapi import (
SensorInfo, SensorInfo,
SensorState, SensorState,
SensorStateClass, SensorStateClass as EsphomeSensorStateClass,
TextSensorInfo, TextSensorInfo,
TextSensorState, TextSensorState,
) )
from aioesphomeapi.model import LastResetType from aioesphomeapi.model import LastResetType
from homeassistant.components.sensor import ( from homeassistant.components.sensor import (
DEVICE_CLASS_TIMESTAMP,
DEVICE_CLASSES, DEVICE_CLASSES,
STATE_CLASS_MEASUREMENT, SensorDeviceClass,
STATE_CLASS_TOTAL_INCREASING,
SensorEntity, SensorEntity,
SensorStateClass,
) )
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
@ -61,11 +60,13 @@ async def async_setup_entry(
# pylint: disable=invalid-overridden-method # pylint: disable=invalid-overridden-method
_STATE_CLASSES: EsphomeEnumMapper[SensorStateClass, str | None] = EsphomeEnumMapper( _STATE_CLASSES: EsphomeEnumMapper[
EsphomeSensorStateClass, SensorStateClass | None
] = EsphomeEnumMapper(
{ {
SensorStateClass.NONE: None, EsphomeSensorStateClass.NONE: None,
SensorStateClass.MEASUREMENT: STATE_CLASS_MEASUREMENT, EsphomeSensorStateClass.MEASUREMENT: SensorStateClass.MEASUREMENT,
SensorStateClass.TOTAL_INCREASING: STATE_CLASS_TOTAL_INCREASING, EsphomeSensorStateClass.TOTAL_INCREASING: SensorStateClass.TOTAL_INCREASING,
} }
) )
@ -85,7 +86,7 @@ class EsphomeSensor(EsphomeEntity[SensorInfo, SensorState], SensorEntity):
return None return None
if self._state.missing_state: if self._state.missing_state:
return None return None
if self.device_class == DEVICE_CLASS_TIMESTAMP: if self.device_class == SensorDeviceClass.TIMESTAMP:
return dt.utc_from_timestamp(self._state.state) return dt.utc_from_timestamp(self._state.state)
return f"{self._state.state:.{self._static_info.accuracy_decimals}f}" return f"{self._state.state:.{self._static_info.accuracy_decimals}f}"
@ -104,18 +105,18 @@ class EsphomeSensor(EsphomeEntity[SensorInfo, SensorState], SensorEntity):
return self._static_info.device_class return self._static_info.device_class
@property @property
def state_class(self) -> str | None: def state_class(self) -> SensorStateClass | None:
"""Return the state class of this entity.""" """Return the state class of this entity."""
if not self._static_info.state_class: if not self._static_info.state_class:
return None return None
state_class = self._static_info.state_class state_class = self._static_info.state_class
reset_type = self._static_info.last_reset_type reset_type = self._static_info.last_reset_type
if ( if (
state_class == SensorStateClass.MEASUREMENT state_class == EsphomeSensorStateClass.MEASUREMENT
and reset_type == LastResetType.AUTO and reset_type == LastResetType.AUTO
): ):
# Legacy, last_reset_type auto was the equivalent to the TOTAL_INCREASING state class # Legacy, last_reset_type auto was the equivalent to the TOTAL_INCREASING state class
return STATE_CLASS_TOTAL_INCREASING return SensorStateClass.TOTAL_INCREASING
return _STATE_CLASSES.from_esphome(self._static_info.state_class) return _STATE_CLASSES.from_esphome(self._static_info.state_class)