Improve person typing (#108218)

This commit is contained in:
Marc Mueller 2024-01-19 01:11:55 +01:00 committed by GitHub
parent 94c8c71ffb
commit a21d5b5858
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,8 +1,9 @@
"""Support for tracking people.""" """Support for tracking people."""
from __future__ import annotations from __future__ import annotations
from collections.abc import Callable
import logging import logging
from typing import Any from typing import Any, Self
import voluptuous as vol import voluptuous as vol
@ -46,10 +47,13 @@ from homeassistant.helpers import (
service, service,
) )
from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.entity_component import EntityComponent
from homeassistant.helpers.event import async_track_state_change_event from homeassistant.helpers.event import (
EventStateChangedData,
async_track_state_change_event,
)
from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.helpers.storage import Store from homeassistant.helpers.storage import Store
from homeassistant.helpers.typing import ConfigType from homeassistant.helpers.typing import ConfigType, EventType
from homeassistant.loader import bind_hass from homeassistant.loader import bind_hass
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -112,7 +116,7 @@ async def async_create_person(
@bind_hass @bind_hass
async def async_add_user_device_tracker( async def async_add_user_device_tracker(
hass: HomeAssistant, user_id: str, device_tracker_entity_id: str hass: HomeAssistant, user_id: str, device_tracker_entity_id: str
): ) -> None:
"""Add a device tracker to a person linked to a user.""" """Add a device tracker to a person linked to a user."""
coll: PersonStorageCollection = hass.data[DOMAIN][1] coll: PersonStorageCollection = hass.data[DOMAIN][1]
@ -187,7 +191,9 @@ UPDATE_FIELDS = {
class PersonStore(Store): class PersonStore(Store):
"""Person storage.""" """Person storage."""
async def _async_migrate_func(self, old_major_version, old_minor_version, old_data): async def _async_migrate_func(
self, old_major_version: int, old_minor_version: int, old_data: dict[str, Any]
) -> dict[str, Any]:
"""Migrate to the new version. """Migrate to the new version.
Migrate storage to use format of collection helper. Migrate storage to use format of collection helper.
@ -281,14 +287,14 @@ class PersonStorageCollection(collection.DictStorageCollection):
"""Return a new updated data object.""" """Return a new updated data object."""
update_data = self.UPDATE_SCHEMA(update_data) update_data = self.UPDATE_SCHEMA(update_data)
user_id = update_data.get(CONF_USER_ID) user_id: str | None = update_data.get(CONF_USER_ID)
if user_id is not None and user_id != item.get(CONF_USER_ID): if user_id is not None and user_id != item.get(CONF_USER_ID):
await self._validate_user_id(user_id) await self._validate_user_id(user_id)
return {**item, **update_data} return {**item, **update_data}
async def _validate_user_id(self, user_id): async def _validate_user_id(self, user_id: str) -> None:
"""Validate the used user_id.""" """Validate the used user_id."""
if await self.hass.auth.async_get_user(user_id) is None: if await self.hass.auth.async_get_user(user_id) is None:
raise ValueError("User does not exist") raise ValueError("User does not exist")
@ -402,32 +408,32 @@ class Person(collection.CollectionEntity, RestoreEntity):
_attr_should_poll = False _attr_should_poll = False
editable: bool editable: bool
def __init__(self, config): def __init__(self, config: dict[str, Any]) -> None:
"""Set up person.""" """Set up person."""
self._config = config self._config = config
self._latitude = None self._latitude: float | None = None
self._longitude = None self._longitude: float | None = None
self._gps_accuracy = None self._gps_accuracy: float | None = None
self._source = None self._source: str | None = None
self._state = None self._state: str | None = None
self._unsub_track_device = None self._unsub_track_device: Callable[[], None] | None = None
@classmethod @classmethod
def from_storage(cls, config: ConfigType): def from_storage(cls, config: ConfigType) -> Self:
"""Return entity instance initialized from storage.""" """Return entity instance initialized from storage."""
person = cls(config) person = cls(config)
person.editable = True person.editable = True
return person return person
@classmethod @classmethod
def from_yaml(cls, config: ConfigType): def from_yaml(cls, config: ConfigType) -> Self:
"""Return entity instance initialized from yaml.""" """Return entity instance initialized from yaml."""
person = cls(config) person = cls(config)
person.editable = False person.editable = False
return person return person
@property @property
def name(self): def name(self) -> str:
"""Return the name of the entity.""" """Return the name of the entity."""
return self._config[CONF_NAME] return self._config[CONF_NAME]
@ -437,14 +443,14 @@ class Person(collection.CollectionEntity, RestoreEntity):
return self._config.get(CONF_PICTURE) return self._config.get(CONF_PICTURE)
@property @property
def state(self): def state(self) -> str | None:
"""Return the state of the person.""" """Return the state of the person."""
return self._state return self._state
@property @property
def extra_state_attributes(self): def extra_state_attributes(self) -> dict[str, Any]:
"""Return the state attributes of the person.""" """Return the state attributes of the person."""
data = {ATTR_EDITABLE: self.editable, ATTR_ID: self.unique_id} data: dict[str, Any] = {ATTR_EDITABLE: self.editable, ATTR_ID: self.unique_id}
if self._latitude is not None: if self._latitude is not None:
data[ATTR_LATITUDE] = self._latitude data[ATTR_LATITUDE] = self._latitude
if self._longitude is not None: if self._longitude is not None:
@ -459,16 +465,16 @@ class Person(collection.CollectionEntity, RestoreEntity):
return data return data
@property @property
def unique_id(self): def unique_id(self) -> str:
"""Return a unique ID for the person.""" """Return a unique ID for the person."""
return self._config[CONF_ID] return self._config[CONF_ID]
@property @property
def device_trackers(self): def device_trackers(self) -> list[str]:
"""Return the device trackers for the person.""" """Return the device trackers for the person."""
return self._config[CONF_DEVICE_TRACKERS] return self._config[CONF_DEVICE_TRACKERS]
async def async_added_to_hass(self): async def async_added_to_hass(self) -> None:
"""Register device trackers.""" """Register device trackers."""
await super().async_added_to_hass() await super().async_added_to_hass()
if state := await self.async_get_last_state(): if state := await self.async_get_last_state():
@ -480,14 +486,14 @@ class Person(collection.CollectionEntity, RestoreEntity):
else: else:
# Wait for hass start to not have race between person # Wait for hass start to not have race between person
# and device trackers finishing setup. # and device trackers finishing setup.
async def person_start_hass(now): async def person_start_hass(_: Event) -> None:
await self.async_update_config(self._config) await self.async_update_config(self._config)
self.hass.bus.async_listen_once( self.hass.bus.async_listen_once(
EVENT_HOMEASSISTANT_START, person_start_hass EVENT_HOMEASSISTANT_START, person_start_hass
) )
async def async_update_config(self, config: ConfigType): async def async_update_config(self, config: ConfigType) -> None:
"""Handle when the config is updated.""" """Handle when the config is updated."""
self._config = config self._config = config
@ -505,12 +511,14 @@ class Person(collection.CollectionEntity, RestoreEntity):
self._update_state() self._update_state()
@callback @callback
def _async_handle_tracker_update(self, event): def _async_handle_tracker_update(
self, event: EventType[EventStateChangedData]
) -> None:
"""Handle the device tracker state changes.""" """Handle the device tracker state changes."""
self._update_state() self._update_state()
@callback @callback
def _update_state(self): def _update_state(self) -> None:
"""Update the state.""" """Update the state."""
latest_non_gps_home = latest_not_home = latest_gps = latest = None latest_non_gps_home = latest_not_home = latest_gps = latest = None
for entity_id in self._config[CONF_DEVICE_TRACKERS]: for entity_id in self._config[CONF_DEVICE_TRACKERS]:
@ -545,7 +553,7 @@ class Person(collection.CollectionEntity, RestoreEntity):
self.async_write_ha_state() self.async_write_ha_state()
@callback @callback
def _parse_source_state(self, state): def _parse_source_state(self, state: State) -> None:
"""Parse source state and set person attributes. """Parse source state and set person attributes.
This is a device tracker state or the restored person state. This is a device tracker state or the restored person state.
@ -570,7 +578,7 @@ def ws_list_person(
) )
def _get_latest(prev: State | None, curr: State): def _get_latest(prev: State | None, curr: State) -> State:
"""Get latest state.""" """Get latest state."""
if prev is None or curr.last_updated > prev.last_updated: if prev is None or curr.last_updated > prev.last_updated:
return curr return curr