Make Store a generic class (#74617)

homeassistant/auth/auth_store.py

@@ -46,7 +46,7 @@ def __init__(self, hass: HomeAssistant) -> None:
         self._users: dict[str, models.User] | None = None
         self._groups: dict[str, models.Group] | None = None
         self._perm_lookup: PermissionLookup | None = None
-        self._store = Store(
+        self._store = Store[dict[str, list[dict[str, Any]]]](
             hass, STORAGE_VERSION, STORAGE_KEY, private=True, atomic_writes=True
         )
         self._lock = asyncio.Lock()
@@ -483,9 +483,10 @@ async def _async_load_task(self) -> None:
                 jwt_key=rt_dict["jwt_key"],
                 last_used_at=last_used_at,
                 last_used_ip=rt_dict.get("last_used_ip"),
-                credential=credentials.get(rt_dict.get("credential_id")),
                 version=rt_dict.get("version"),
             )
+            if "credential_id" in rt_dict:
+                token.credential = credentials.get(rt_dict["credential_id"])
             users[rt_dict["user_id"]].refresh_tokens[token.id] = token
 
         self._groups = groups

homeassistant/auth/mfa_modules/notify.py

@@ -7,7 +7,7 @@
 import asyncio
 from collections import OrderedDict
 import logging
-from typing import Any
+from typing import Any, cast
 
 import attr
 import voluptuous as vol
@@ -100,7 +100,7 @@ def __init__(self, hass: HomeAssistant, config: dict[str, Any]) -> None:
         """Initialize the user data store."""
         super().__init__(hass, config)
         self._user_settings: _UsersDict | None = None
-        self._user_store = Store(
+        self._user_store = Store[dict[str, dict[str, Any]]](
             hass, STORAGE_VERSION, STORAGE_KEY, private=True, atomic_writes=True
         )
         self._include = config.get(CONF_INCLUDE, [])
@@ -119,10 +119,8 @@ async def _async_load(self) -> None:
             if self._user_settings is not None:
                 return
 
-            if (data := await self._user_store.async_load()) is None or not isinstance(
-                data, dict
-            ):
-                data = {STORAGE_USERS: {}}
+            if (data := await self._user_store.async_load()) is None:
+                data = cast(dict[str, dict[str, Any]], {STORAGE_USERS: {}})
 
             self._user_settings = {
                 user_id: NotifySetting(**setting)

homeassistant/auth/mfa_modules/totp.py

@@ -3,7 +3,7 @@
 
 import asyncio
 from io import BytesIO
-from typing import Any
+from typing import Any, cast
 
 import voluptuous as vol
 
@@ -77,7 +77,7 @@ def __init__(self, hass: HomeAssistant, config: dict[str, Any]) -> None:
         """Initialize the user data store."""
         super().__init__(hass, config)
         self._users: dict[str, str] | None = None
-        self._user_store = Store(
+        self._user_store = Store[dict[str, dict[str, str]]](
             hass, STORAGE_VERSION, STORAGE_KEY, private=True, atomic_writes=True
         )
         self._init_lock = asyncio.Lock()
@@ -93,16 +93,14 @@ async def _async_load(self) -> None:
             if self._users is not None:
                 return
 
-            if (data := await self._user_store.async_load()) is None or not isinstance(
-                data, dict
-            ):
-                data = {STORAGE_USERS: {}}
+            if (data := await self._user_store.async_load()) is None:
+                data = cast(dict[str, dict[str, str]], {STORAGE_USERS: {}})
 
             self._users = data.get(STORAGE_USERS, {})
 
     async def _async_save(self) -> None:
         """Save data."""
-        await self._user_store.async_save({STORAGE_USERS: self._users})
+        await self._user_store.async_save({STORAGE_USERS: self._users or {}})
 
     def _add_ota_secret(self, user_id: str, secret: str | None = None) -> str:
         """Create a ota_secret for user."""

homeassistant/auth/providers/homeassistant.py

@@ -61,10 +61,10 @@ class Data:
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the user data store."""
         self.hass = hass
-        self._store = Store(
+        self._store = Store[dict[str, list[dict[str, str]]]](
             hass, STORAGE_VERSION, STORAGE_KEY, private=True, atomic_writes=True
         )
-        self._data: dict[str, Any] | None = None
+        self._data: dict[str, list[dict[str, str]]] | None = None
         # Legacy mode will allow usernames to start/end with whitespace
         # and will compare usernames case-insensitive.
         # Remove in 2020 or when we launch 1.0.
@@ -80,10 +80,8 @@ def normalize_username(self, username: str) -> str:
 
     async def async_load(self) -> None:
         """Load stored data."""
-        if (data := await self._store.async_load()) is None or not isinstance(
-            data, dict
-        ):
-            data = {"users": []}
+        if (data := await self._store.async_load()) is None:
+            data = cast(dict[str, list[dict[str, str]]], {"users": []})
 
         seen: set[str] = set()
 
@@ -123,7 +121,8 @@ async def async_load(self) -> None:
     @property
     def users(self) -> list[dict[str, str]]:
         """Return users."""
-        return self._data["users"]  # type: ignore[index,no-any-return]
+        assert self._data is not None
+        return self._data["users"]
 
     def validate_login(self, username: str, password: str) -> None:
         """Validate a username and password.

homeassistant/components/almond/__init__.py

@@ -5,7 +5,7 @@
 from datetime import timedelta
 import logging
 import time
-from typing import Optional, cast
+from typing import Any
 
 from aiohttp import ClientError, ClientSession
 import async_timeout
@@ -167,8 +167,8 @@ async def _configure_almond_for_ha(
         return
 
     _LOGGER.debug("Configuring Almond to connect to Home Assistant at %s", hass_url)
-    store = storage.Store(hass, STORAGE_VERSION, STORAGE_KEY)
-    data = cast(Optional[dict], await store.async_load())
+    store = storage.Store[dict[str, Any]](hass, STORAGE_VERSION, STORAGE_KEY)
+    data = await store.async_load()
 
     if data is None:
         data = {}

homeassistant/components/ambiclimate/climate.py

@@ -64,7 +64,7 @@ async def async_setup_entry(
     """Set up the Ambiclimate device from config entry."""
     config = entry.data
     websession = async_get_clientsession(hass)
-    store = Store(hass, STORAGE_VERSION, STORAGE_KEY)
+    store = Store[dict[str, Any]](hass, STORAGE_VERSION, STORAGE_KEY)
     token_info = await store.async_load()
 
     oauth = ambiclimate.AmbiclimateOAuth(

homeassistant/components/analytics/analytics.py

@@ -1,6 +1,6 @@
 """Analytics helper class for the analytics integration."""
 import asyncio
-from typing import cast
+from typing import Any
 import uuid
 
 import aiohttp
@@ -66,12 +66,12 @@ def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the Analytics class."""
         self.hass: HomeAssistant = hass
         self.session = async_get_clientsession(hass)
-        self._data: dict = {
+        self._data: dict[str, Any] = {
             ATTR_PREFERENCES: {},
             ATTR_ONBOARDED: False,
             ATTR_UUID: None,
         }
-        self._store = Store(hass, STORAGE_VERSION, STORAGE_KEY)
+        self._store = Store[dict[str, Any]](hass, STORAGE_VERSION, STORAGE_KEY)
 
     @property
     def preferences(self) -> dict:
@@ -109,7 +109,7 @@ def supervisor(self) -> bool:
 
     async def load(self) -> None:
         """Load preferences."""
-        stored = cast(dict, await self._store.async_load())
+        stored = await self._store.async_load()
         if stored:
             self._data = stored
 

homeassistant/components/camera/prefs.py

@@ -36,14 +36,14 @@ class CameraPreferences:
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize camera prefs."""
         self._hass = hass
-        self._store = Store(hass, STORAGE_VERSION, STORAGE_KEY)
+        self._store = Store[dict[str, dict[str, bool]]](
+            hass, STORAGE_VERSION, STORAGE_KEY
+        )
         self._prefs: dict[str, dict[str, bool]] | None = None
 
     async def async_initialize(self) -> None:
         """Finish initializing the preferences."""
-        if (prefs := await self._store.async_load()) is None or not isinstance(
-            prefs, dict
-        ):
+        if (prefs := await self._store.async_load()) is None:
             prefs = {}
 
         self._prefs = prefs

homeassistant/components/energy/data.py

@@ -4,7 +4,7 @@
 import asyncio
 from collections import Counter
 from collections.abc import Awaitable, Callable
-from typing import Literal, Optional, TypedDict, Union, cast
+from typing import Literal, TypedDict, Union
 
 import voluptuous as vol
 
@@ -263,13 +263,15 @@ class EnergyManager:
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize energy manager."""
         self._hass = hass
-        self._store = storage.Store(hass, STORAGE_VERSION, STORAGE_KEY)
+        self._store = storage.Store[EnergyPreferences](
+            hass, STORAGE_VERSION, STORAGE_KEY
+        )
         self.data: EnergyPreferences | None = None
         self._update_listeners: list[Callable[[], Awaitable]] = []
 
     async def async_initialize(self) -> None:
         """Initialize the energy integration."""
-        self.data = cast(Optional[EnergyPreferences], await self._store.async_load())
+        self.data = await self._store.async_load()
 
     @staticmethod
     def default_preferences() -> EnergyPreferences:
@@ -294,7 +296,7 @@ async def async_update(self, update: EnergyPreferencesUpdate) -> None:
                 data[key] = update[key]  # type: ignore[literal-required]
 
         self.data = data
-        self._store.async_delay_save(lambda: cast(dict, self.data), 60)
+        self._store.async_delay_save(lambda: data, 60)
 
         if not self._update_listeners:
             return

homeassistant/components/hassio/__init__.py

@@ -533,12 +533,10 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:  # noqa:
     if not await hassio.is_connected():
         _LOGGER.warning("Not connected with the supervisor / system too busy!")
 
-    store = Store(hass, STORAGE_VERSION, STORAGE_KEY)
+    store = Store[dict[str, str]](hass, STORAGE_VERSION, STORAGE_KEY)
     if (data := await store.async_load()) is None:
         data = {}
 
-    assert isinstance(data, dict)
-
     refresh_token = None
     if "hassio_user" in data:
         user = await hass.auth.async_get_user(data["hassio_user"])

homeassistant/components/homekit_controller/storage.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, TypedDict, cast
+from typing import Any, TypedDict
 
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.storage import Store
@@ -46,7 +46,9 @@ class EntityMapStorage:
     def __init__(self, hass: HomeAssistant) -> None:
         """Create a new entity map store."""
         self.hass = hass
-        self.store = Store(hass, ENTITY_MAP_STORAGE_VERSION, ENTITY_MAP_STORAGE_KEY)
+        self.store = Store[StorageLayout](
+            hass, ENTITY_MAP_STORAGE_VERSION, ENTITY_MAP_STORAGE_KEY
+        )
         self.storage_data: dict[str, Pairing] = {}
 
     async def async_initialize(self) -> None:
@@ -55,8 +57,7 @@ async def async_initialize(self) -> None:
             # There is no cached data about HomeKit devices yet
             return
 
-        storage = cast(StorageLayout, raw_storage)
-        self.storage_data = storage.get("pairings", {})
+        self.storage_data = raw_storage.get("pairings", {})
 
     def get_map(self, homekit_id: str) -> Pairing | None:
         """Get a pairing cache item."""
@@ -87,6 +88,6 @@ def _async_schedule_save(self) -> None:
         self.store.async_delay_save(self._data_to_save, ENTITY_MAP_SAVE_DELAY)
 
     @callback
-    def _data_to_save(self) -> dict[str, Any]:
+    def _data_to_save(self) -> StorageLayout:
         """Return data of entity map to store in a file."""
-        return {"pairings": self.storage_data}
+        return StorageLayout(pairings=self.storage_data)

homeassistant/components/http/__init__.py

@@ -7,7 +7,7 @@
 import os
 import ssl
 from tempfile import NamedTemporaryFile
-from typing import Any, Final, Optional, TypedDict, Union, cast
+from typing import Any, Final, TypedDict, Union, cast
 
 from aiohttp import web
 from aiohttp.typedefs import StrOrURL
@@ -125,10 +125,10 @@ class ConfData(TypedDict, total=False):
 
 
 @bind_hass
-async def async_get_last_config(hass: HomeAssistant) -> dict | None:
+async def async_get_last_config(hass: HomeAssistant) -> dict[str, Any] | None:
     """Return the last known working config."""
-    store = storage.Store(hass, STORAGE_VERSION, STORAGE_KEY)
-    return cast(Optional[dict], await store.async_load())
+    store = storage.Store[dict[str, Any]](hass, STORAGE_VERSION, STORAGE_KEY)
+    return await store.async_load()
 
 
 class ApiConfig:
@@ -475,7 +475,9 @@ async def start_http_server_and_save_config(
     await server.start()
 
     # If we are set up successful, we store the HTTP settings for safe mode.
-    store = storage.Store(hass, STORAGE_VERSION, STORAGE_KEY)
+    store: storage.Store[dict[str, Any]] = storage.Store(
+        hass, STORAGE_VERSION, STORAGE_KEY
+    )
 
     if CONF_TRUSTED_PROXIES in conf:
         conf[CONF_TRUSTED_PROXIES] = [

homeassistant/components/http/auth.py

@@ -6,7 +6,7 @@
 from ipaddress import ip_address
 import logging
 import secrets
-from typing import Final
+from typing import Any, Final
 
 from aiohttp import hdrs
 from aiohttp.web import Application, Request, StreamResponse, middleware
@@ -118,8 +118,8 @@ def async_user_not_allowed_do_auth(
 
 async def async_setup_auth(hass: HomeAssistant, app: Application) -> None:
     """Create auth middleware for the app."""
-    store = Store(hass, STORAGE_VERSION, STORAGE_KEY)
-    if (data := await store.async_load()) is None or not isinstance(data, dict):
+    store = Store[dict[str, Any]](hass, STORAGE_VERSION, STORAGE_KEY)
+    if (data := await store.async_load()) is None:
         data = {}
 
     refresh_token = None

homeassistant/components/mobile_app/__init__.py

@@ -1,5 +1,6 @@
 """Integrates Native Apps to Home Assistant."""
 from contextlib import suppress
+from typing import Any
 
 from homeassistant.components import cloud, notify as hass_notify
 from homeassistant.components.webhook import (
@@ -38,7 +39,7 @@
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the mobile app component."""
-    store = Store(hass, STORAGE_VERSION, STORAGE_KEY)
+    store = Store[dict[str, Any]](hass, STORAGE_VERSION, STORAGE_KEY)
     if (app_config := await store.async_load()) is None or not isinstance(
         app_config, dict
     ):