evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Improve registry store data typing (#115066)

Browse source
This commit is contained in:
Marc Mueller 2024-04-07 02:34:49 +02:00 • committed by GitHub
parent a0936902c2
commit cb9352110c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 104 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
homeassistant/helpers/area_registry.py
View file

@ -26,6 +26,24 @@ STORAGE_VERSION_MAJOR = 1
STORAGE_VERSION_MINOR = 6
class _AreaStoreData(TypedDict):
"""Data type for individual area. Used in AreasRegistryStoreData."""
aliases: list[str]
floor_id: str | None
icon: str | None
id: str
labels: list[str]
name: str
picture: str | None
class AreasRegistryStoreData(TypedDict):
"""Store data type for AreaRegistry."""
areas: list[_AreaStoreData]
class EventAreaRegistryUpdatedData(TypedDict):
"""EventAreaRegistryUpdated data."""
@ -45,7 +63,7 @@ class AreaEntry(NormalizedNameBaseRegistryEntry):
picture: str | None
class AreaRegistryStore(Store[dict[str, list[dict[str, Any]]]]):
class AreaRegistryStore(Store[AreasRegistryStoreData]):
"""Store area registry data."""
async def _async_migrate_func(
@ -53,7 +71,7 @@ class AreaRegistryStore(Store[dict[str, list[dict[str, Any]]]]):
old_major_version: int,
old_minor_version: int,
old_data: dict[str, list[dict[str, Any]]],
) -> dict[str, Any]:
) -> AreasRegistryStoreData:
"""Migrate to the new version."""
if old_major_version < 2:
if old_minor_version < 2:
@ -84,7 +102,7 @@ class AreaRegistryStore(Store[dict[str, list[dict[str, Any]]]]):
if old_major_version > 1:
raise NotImplementedError
return old_data
return old_data # type: ignore[return-value]
class AreaRegistryItems(NormalizedNameBaseRegistryItems[AreaEntry]):
@ -126,7 +144,7 @@ class AreaRegistryItems(NormalizedNameBaseRegistryItems[AreaEntry]):
return [data[key] for key in self._floors_index.get(floor, ())]
class AreaRegistry(BaseRegistry):
class AreaRegistry(BaseRegistry[AreasRegistryStoreData]):
"""Class to hold a registry of areas."""
areas: AreaRegistryItems
@ -314,24 +332,22 @@ class AreaRegistry(BaseRegistry):
self._area_data = areas.data
@callback
def _data_to_save(self) -> dict[str, list[dict[str, Any]]]:
def _data_to_save(self) -> AreasRegistryStoreData:
"""Return data of area registry to store in a file."""
data = {}
data["areas"] = [
{
"aliases": list(entry.aliases),
"floor_id": entry.floor_id,
"icon": entry.icon,
"id": entry.id,
"labels": list(entry.labels),
"name": entry.name,
"picture": entry.picture,
}
for entry in self.areas.values()
]
return data
return {
"areas": [
{
"aliases": list(entry.aliases),
"floor_id": entry.floor_id,
"icon": entry.icon,
"id": entry.id,
"labels": list(entry.labels),
"name": entry.name,
"picture": entry.picture,
}
for entry in self.areas.values()
]
}
def _generate_area_id(self, name: str) -> str:
"""Generate area ID."""

20
homeassistant/helpers/category_registry.py
View file

@ -20,6 +20,20 @@ STORAGE_KEY = "core.category_registry"
STORAGE_VERSION_MAJOR = 1
class _CategoryStoreData(TypedDict):
"""Data type for individual category. Used in CategoryRegistryStoreData."""
category_id: str
icon: str | None
name: str
class CategoryRegistryStoreData(TypedDict):
"""Store data type for CategoryRegistry."""
categories: dict[str, list[_CategoryStoreData]]
class EventCategoryRegistryUpdatedData(TypedDict):
"""Event data for when the category registry is updated."""
@ -40,14 +54,14 @@ class CategoryEntry:
name: str
class CategoryRegistry(BaseRegistry):
class CategoryRegistry(BaseRegistry[CategoryRegistryStoreData]):
"""Class to hold a registry of categories by scope."""
def __init__(self, hass: HomeAssistant) -> None:
"""Initialize the category registry."""
self.hass = hass
self.categories: dict[str, dict[str, CategoryEntry]] = {}
self._store: Store[dict[str, dict[str, list[dict[str, str]]]]] = Store(
self._store = Store(
hass,
STORAGE_VERSION_MAJOR,
STORAGE_KEY,
@ -167,7 +181,7 @@ class CategoryRegistry(BaseRegistry):
self.categories = category_entries
@callback
def _data_to_save(self) -> dict[str, dict[str, list[dict[str, str | None]]]]:
def _data_to_save(self) -> CategoryRegistryStoreData:
"""Return data of category registry to store in a file."""
return {
"categories": {

2
homeassistant/helpers/device_registry.py
View file

@ -551,7 +551,7 @@ class ActiveDeviceRegistryItems(DeviceRegistryItems[DeviceEntry]):
]
class DeviceRegistry(BaseRegistry):
class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
"""Class to hold a registry of devices."""
devices: ActiveDeviceRegistryItems

41
homeassistant/helpers/floor_registry.py
View file

@ -5,7 +5,7 @@ from __future__ import annotations
from collections.abc import Iterable
import dataclasses
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal, TypedDict, cast
from typing import Literal, TypedDict, cast
from homeassistant.core import Event, HomeAssistant, callback
from homeassistant.util import slugify
@ -25,6 +25,22 @@ STORAGE_KEY = "core.floor_registry"
STORAGE_VERSION_MAJOR = 1
class _FloorStoreData(TypedDict):
"""Data type for individual floor. Used in FloorRegistryStoreData."""
aliases: list[str]
floor_id: str
icon: str | None
level: int | None
name: str
class FloorRegistryStoreData(TypedDict):
"""Store data type for FloorRegistry."""
floors: list[_FloorStoreData]
class EventFloorRegistryUpdatedData(TypedDict):
"""Event data for when the floor registry is updated."""
@ -45,7 +61,7 @@ class FloorEntry(NormalizedNameBaseRegistryEntry):
level: int | None = None
class FloorRegistry(BaseRegistry):
class FloorRegistry(BaseRegistry[FloorRegistryStoreData]):
"""Class to hold a registry of floors."""
floors: NormalizedNameBaseRegistryItems[FloorEntry]
@ -54,13 +70,11 @@ class FloorRegistry(BaseRegistry):
def __init__(self, hass: HomeAssistant) -> None:
"""Initialize the floor registry."""
self.hass = hass
self._store: Store[dict[str, list[dict[str, str | int | list[str] | None]]]] = (
Store(
hass,
STORAGE_VERSION_MAJOR,
STORAGE_KEY,
atomic_writes=True,
)
self._store = Store(
hass,
STORAGE_VERSION_MAJOR,
STORAGE_KEY,
atomic_writes=True,
)
@callback
@ -190,13 +204,6 @@ class FloorRegistry(BaseRegistry):
if data is not None:
for floor in data["floors"]:
if TYPE_CHECKING:
assert isinstance(floor["aliases"], list)
assert isinstance(floor["icon"], str)
assert isinstance(floor["level"], int)
assert isinstance(floor["name"], str)
assert isinstance(floor["floor_id"], str)
normalized_name = normalize_name(floor["name"])
floors[floor["floor_id"]] = FloorEntry(
aliases=set(floor["aliases"]),
@ -211,7 +218,7 @@ class FloorRegistry(BaseRegistry):
self._floor_data = floors.data
@callback
def _data_to_save(self) -> dict[str, list[dict[str, str | int | list[str] | None]]]:
def _data_to_save(self) -> FloorRegistryStoreData:
"""Return data of floor registry to store in a file."""
return {
"floors": [

26
homeassistant/helpers/label_registry.py
View file

@ -25,6 +25,22 @@ STORAGE_KEY = "core.label_registry"
STORAGE_VERSION_MAJOR = 1
class _LabelStoreData(TypedDict):
"""Data type for individual label. Used in LabelRegistryStoreData."""
color: str | None
description: str | None
icon: str | None
label_id: str
name: str
class LabelRegistryStoreData(TypedDict):
"""Store data type for LabelRegistry."""
labels: list[_LabelStoreData]
class EventLabelRegistryUpdatedData(TypedDict):
"""Event data for when the label registry is updated."""
@ -45,7 +61,7 @@ class LabelEntry(NormalizedNameBaseRegistryEntry):
icon: str | None = None
class LabelRegistry(BaseRegistry):
class LabelRegistry(BaseRegistry[LabelRegistryStoreData]):
"""Class to hold a registry of labels."""
labels: NormalizedNameBaseRegistryItems[LabelEntry]
@ -54,7 +70,7 @@ class LabelRegistry(BaseRegistry):
def __init__(self, hass: HomeAssistant) -> None:
"""Initialize the label registry."""
self.hass = hass
self._store: Store[dict[str, list[dict[str, str | None]]]] = Store(
self._store = Store(
hass,
STORAGE_VERSION_MAJOR,
STORAGE_KEY,
@ -189,10 +205,6 @@ class LabelRegistry(BaseRegistry):
if data is not None:
for label in data["labels"]:
# Check if the necessary keys are present
if label["label_id"] is None or label["name"] is None:
continue
normalized_name = normalize_name(label["name"])
labels[label["label_id"]] = LabelEntry(
color=label["color"],
@ -207,7 +219,7 @@ class LabelRegistry(BaseRegistry):
self._label_data = labels.data
@callback
def _data_to_save(self) -> dict[str, list[dict[str, str | None]]]:
def _data_to_save(self) -> LabelRegistryStoreData:
"""Return data of label registry to store in a file."""
return {
"labels": [

11
homeassistant/helpers/registry.py
View file

@ -4,8 +4,8 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from collections import UserDict
from collections.abc import ValuesView
from typing import TYPE_CHECKING, Any, Literal, TypeVar
from collections.abc import Mapping, Sequence, ValuesView
from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar
from homeassistant.core import CoreState, HomeAssistant, callback
@ -17,6 +17,7 @@ SAVE_DELAY_LONG = 180
_DataT = TypeVar("_DataT")
_StoreDataT = TypeVar("_StoreDataT", bound=Mapping[str, Any] | Sequence[Any])
class BaseRegistryItems(UserDict[str, _DataT], ABC):
@ -64,11 +65,11 @@ class BaseRegistryItems(UserDict[str, _DataT], ABC):
super().__delitem__(key)
class BaseRegistry(ABC):
class BaseRegistry(ABC, Generic[_StoreDataT]):
"""Class to implement a registry."""
hass: HomeAssistant
_store: Store
_store: Store[_StoreDataT]
@callback
def async_schedule_save(self) -> None:
@ -80,5 +81,5 @@ class BaseRegistry(ABC):
@callback
@abstractmethod
def _data_to_save(self) -> dict[str, Any]:
def _data_to_save(self) -> _StoreDataT:
"""Return data of registry to store in a file."""