457 lines
18 KiB
Python
457 lines
18 KiB
Python
import asyncio
|
|
from collections.abc import Callable, Coroutine
|
|
from functools import partial
|
|
import json
|
|
from json import JSONDecodeError
|
|
import logging
|
|
import os
|
|
import shutil
|
|
from typing import Any, NotRequired, TypedDict, cast
|
|
|
|
import aiohttp
|
|
from aiohttp import ClientError
|
|
from awesomeversion import AwesomeVersion
|
|
from homeassistant.core import HomeAssistant
|
|
from homeassistant.helpers.aiohttp_client import async_get_clientsession
|
|
from homeassistant.helpers.storage import STORAGE_DIR
|
|
from homeassistant.loader import async_get_integration
|
|
|
|
from custom_components.powercalc.const import API_URL, BUILT_IN_LIBRARY_DIR, DOMAIN
|
|
from custom_components.powercalc.helpers import async_cache, clear_async_cache
|
|
from custom_components.powercalc.power_profile.error import LibraryLoadingError, ProfileDownloadError
|
|
from custom_components.powercalc.power_profile.loader.protocol import Loader
|
|
from custom_components.powercalc.power_profile.power_profile import DeviceType, DiscoveryBy
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
ENDPOINT_LIBRARY = f"{API_URL}/library"
|
|
ENDPOINT_DOWNLOAD = f"{API_URL}/download"
|
|
|
|
TIMEOUT_SECONDS = 30
|
|
|
|
|
|
class LibraryModel(TypedDict):
|
|
id: str
|
|
name: NotRequired[str]
|
|
aliases: NotRequired[list[str]]
|
|
legacy_ids: NotRequired[list[str]]
|
|
hash: str
|
|
device_type: NotRequired[DeviceType]
|
|
discovery_by: NotRequired[DiscoveryBy]
|
|
min_version: NotRequired[str]
|
|
|
|
|
|
class LibraryManufacturer(TypedDict):
|
|
name: str
|
|
dir_name: str
|
|
aliases: NotRequired[list[str]]
|
|
models: list[LibraryModel]
|
|
|
|
|
|
class RemoteLoader(Loader):
|
|
retry_timeout = 3
|
|
|
|
def __init__(self, hass: HomeAssistant) -> None:
|
|
self.hass = hass
|
|
self.library_contents: dict = {}
|
|
self.model_infos: dict[str, LibraryModel] = {}
|
|
self.manufacturer_models: dict[str, list[LibraryModel]] = {}
|
|
self.model_lookup: dict[str, dict[str, list[LibraryModel]]] = {}
|
|
self.manufacturer_lookup: dict[str, set[str]] = {}
|
|
self.profile_hashes: dict[str, str] = {}
|
|
|
|
async def initialize(self) -> None:
|
|
"""Initialize the loader."""
|
|
|
|
integration = await async_get_integration(self.hass, DOMAIN)
|
|
powercalc_version = AwesomeVersion(str(integration.version))
|
|
|
|
self._clear_caches()
|
|
self.library_contents = await self.load_library_json()
|
|
self.profile_hashes = await self.hass.async_add_executor_job(self._load_profile_hashes)
|
|
|
|
self.model_infos.clear()
|
|
self.model_lookup.clear()
|
|
self.manufacturer_models.clear()
|
|
self.manufacturer_lookup.clear()
|
|
|
|
manufacturers: list[LibraryManufacturer] = self.library_contents.get("manufacturers", [])
|
|
|
|
for manufacturer in manufacturers:
|
|
manufacturer_name = str(manufacturer.get("dir_name"))
|
|
models: list[LibraryModel] = manufacturer.get("models", []) or []
|
|
|
|
# manufacturer alias map (alias -> {canonical manufacturer_name})
|
|
self.manufacturer_lookup.setdefault(manufacturer_name.lower(), set()).add(manufacturer_name)
|
|
for alias in manufacturer.get("aliases", []) or []:
|
|
self.manufacturer_lookup.setdefault(str(alias).lower(), set()).add(manufacturer_name)
|
|
|
|
# per-manufacturer model lookup
|
|
kept_models: list[LibraryModel] = []
|
|
lookup: dict[str, list[LibraryModel]] = {}
|
|
|
|
for model in models:
|
|
min_version = model.get("min_version")
|
|
model_id = str(model.get("id"))
|
|
model_id_lower = model_id.lower()
|
|
|
|
self.model_infos[f"{manufacturer_name}/{model_id!s}"] = model
|
|
|
|
if min_version and powercalc_version < AwesomeVersion(min_version):
|
|
_LOGGER.debug(
|
|
"Skipping model %s/%s as it requires powercalc version %s (current: %s)",
|
|
manufacturer_name,
|
|
model_id,
|
|
min_version,
|
|
powercalc_version,
|
|
)
|
|
continue
|
|
|
|
kept_models.append(model)
|
|
|
|
# Exact id bucket first (highest priority)
|
|
bucket = lookup.setdefault(model_id_lower, [])
|
|
bucket.insert(0, model)
|
|
|
|
# Alias buckets afterwards (lower priority)
|
|
for alias in model.get("aliases", []) or []:
|
|
alias_lower = str(alias).lower()
|
|
if alias_lower == model_id_lower:
|
|
continue
|
|
# Append to the end to ensure aliased models are always last
|
|
lookup.setdefault(alias_lower, []).append(model)
|
|
|
|
self.manufacturer_models[manufacturer_name] = kept_models
|
|
self.model_lookup[manufacturer_name] = lookup
|
|
|
|
def _clear_caches(self) -> None:
|
|
"""Clear cached lookups backed by mutable library state."""
|
|
clear_async_cache(self.get_manufacturer_listing)
|
|
clear_async_cache(self.find_manufacturers)
|
|
clear_async_cache(self.get_model_listing)
|
|
clear_async_cache(self.find_model)
|
|
clear_async_cache(self.find_model_migration)
|
|
clear_async_cache(self.load_model)
|
|
|
|
async def load_library_json(self) -> dict[str, Any]:
|
|
"""Load library.json file"""
|
|
|
|
local_path = self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, "library.json")
|
|
|
|
def _load_local_library_json() -> dict[str, Any]:
|
|
"""Load library.json file from local storage"""
|
|
if not os.path.exists(local_path):
|
|
raise ProfileDownloadError("Local library.json file not found")
|
|
with open(local_path) as f:
|
|
return cast(dict[str, Any], json.load(f))
|
|
|
|
async def _download_remote_library_json() -> dict[str, Any] | None:
|
|
"""
|
|
Download library.json from Github.
|
|
On success, save it to local storage as a fallback for internet connection issues.
|
|
"""
|
|
_LOGGER.debug("Loading library.json from github")
|
|
|
|
session = async_get_clientsession(self.hass)
|
|
|
|
try:
|
|
async with asyncio.timeout(TIMEOUT_SECONDS), session.get(ENDPOINT_LIBRARY) as resp:
|
|
if resp.status != 200:
|
|
raise ProfileDownloadError(
|
|
f"Failed to download library.json, unexpected status code: {resp.status}",
|
|
)
|
|
|
|
data = await resp.read()
|
|
|
|
except (TimeoutError, ClientError) as err:
|
|
raise ProfileDownloadError(f"Failed to download library.json: {err}") from err
|
|
|
|
def _save_to_local_storage(data: bytes) -> None:
|
|
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
|
with open(local_path, "wb") as f:
|
|
f.write(data)
|
|
|
|
await self.hass.async_add_executor_job(_save_to_local_storage, data)
|
|
|
|
return cast(dict[str, Any], json.loads(data))
|
|
|
|
try:
|
|
return cast(dict[str, Any], await self.download_with_retry(_download_remote_library_json))
|
|
except ProfileDownloadError:
|
|
_LOGGER.debug("Failed to download library.json, falling back to local copy")
|
|
return await self.hass.async_add_executor_job(_load_local_library_json)
|
|
|
|
@async_cache
|
|
async def get_manufacturer_listing(
|
|
self,
|
|
device_types: set[DeviceType] | None,
|
|
discovery_by: DiscoveryBy | None = None,
|
|
) -> set[tuple[str, str]]:
|
|
"""Get listing of available manufacturers."""
|
|
|
|
return {
|
|
(manufacturer["dir_name"], manufacturer["full_name"])
|
|
for manufacturer in self.library_contents.get("manufacturers", [])
|
|
if any(
|
|
self._model_matches_filters(model, device_types, discovery_by)
|
|
for model in manufacturer.get("models", [])
|
|
)
|
|
}
|
|
|
|
@async_cache
|
|
async def find_manufacturers(self, search: str) -> set[str]:
|
|
"""Find the manufacturer in the library."""
|
|
return self.manufacturer_lookup.get(search.lower(), set())
|
|
|
|
@async_cache
|
|
async def get_model_listing(
|
|
self,
|
|
manufacturer: str,
|
|
device_types: set[DeviceType] | None,
|
|
discovery_by: DiscoveryBy | None = None,
|
|
) -> set[tuple[str, str]]:
|
|
"""Get listing of available models and display names for a given manufacturer."""
|
|
models = self.manufacturer_models.get(manufacturer)
|
|
if not models:
|
|
return set()
|
|
|
|
return {
|
|
(model["id"], str(model.get("name") or model["id"]))
|
|
for model in self.manufacturer_models.get(manufacturer, [])
|
|
if self._model_matches_filters(model, device_types, discovery_by)
|
|
}
|
|
|
|
@staticmethod
|
|
def _model_matches_filters(
|
|
model: LibraryModel,
|
|
device_types: set[DeviceType] | None,
|
|
discovery_by: DiscoveryBy | None,
|
|
) -> bool:
|
|
model_device_type = DeviceType(model.get("device_type", DeviceType.LIGHT))
|
|
if device_types and model_device_type not in device_types:
|
|
return False
|
|
|
|
model_discovery_by = DiscoveryBy(model.get("discovery_by", DiscoveryBy.ENTITY))
|
|
return not discovery_by or model_discovery_by == discovery_by
|
|
|
|
@async_cache
|
|
async def find_model(self, manufacturer: str, search: set[str]) -> list[str]:
|
|
"""Find matching model IDs in the library."""
|
|
models = self.model_lookup.get(manufacturer, {})
|
|
return [
|
|
model["id"]
|
|
for phrase in search
|
|
if (phrase_lower := phrase.lower()) in models
|
|
for model in models[phrase_lower]
|
|
]
|
|
|
|
@async_cache
|
|
async def find_model_migration(self, manufacturer: str, model: str) -> str | None:
|
|
"""Find the canonical model id for a legacy profile id."""
|
|
model_lower = model.lower()
|
|
matches = {
|
|
str(model_data.get("id"))
|
|
for manufacturer_data in self.library_contents.get("manufacturers", [])
|
|
if str(manufacturer_data.get("dir_name", "")).lower() == manufacturer
|
|
for model_data in manufacturer_data.get("models", []) or []
|
|
if model_lower in {str(legacy_id).lower() for legacy_id in model_data.get("legacy_ids", []) or []}
|
|
}
|
|
|
|
if len(matches) != 1:
|
|
return None
|
|
|
|
return next(iter(matches))
|
|
|
|
@async_cache
|
|
async def load_model(
|
|
self,
|
|
manufacturer: str,
|
|
model: str,
|
|
force_update: bool = False,
|
|
retry_count: int = 0,
|
|
) -> tuple[dict, str] | None:
|
|
"""Load a model, downloading it if necessary, with retry logic."""
|
|
model_info = self._get_library_model(manufacturer, model)
|
|
storage_path = self.get_storage_path(manufacturer, model)
|
|
model_path = os.path.join(storage_path, "model.json")
|
|
|
|
if await self._needs_update(model_info, manufacturer, model, model_path, force_update):
|
|
await self._download_profile_with_retry(manufacturer, model, storage_path, model_path)
|
|
|
|
try:
|
|
json_data = await self._load_model_json(model_path)
|
|
except JSONDecodeError as e:
|
|
return await self._handle_json_decode_error(e, manufacturer, model, retry_count)
|
|
|
|
return json_data, storage_path
|
|
|
|
def _get_library_model(self, manufacturer: str, model: str) -> LibraryModel:
|
|
"""Retrieve model info, or raise an error if not found."""
|
|
model_info = self.model_infos.get(f"{manufacturer}/{model}")
|
|
if not model_info:
|
|
raise LibraryLoadingError(f"Model not found in library: {manufacturer}/{model}")
|
|
return model_info
|
|
|
|
async def _needs_update(
|
|
self,
|
|
model_info: LibraryModel,
|
|
manufacturer: str,
|
|
model: str,
|
|
model_path: str,
|
|
force_update: bool,
|
|
) -> bool:
|
|
"""Check if the model needs to be updated."""
|
|
if force_update:
|
|
return True
|
|
|
|
path_exists = await self.hass.async_add_executor_job(os.path.exists, model_path)
|
|
if not path_exists:
|
|
return True
|
|
|
|
existing_hash = self.profile_hashes.get(f"{manufacturer}/{model}")
|
|
new_hash = model_info.get("hash")
|
|
return existing_hash != new_hash
|
|
|
|
async def _download_profile_with_retry(
|
|
self,
|
|
manufacturer: str,
|
|
model: str,
|
|
storage_path: str,
|
|
model_path: str,
|
|
) -> None:
|
|
"""Attempt to download the profile, with retry logic and error handling."""
|
|
try:
|
|
model_info = self._get_library_model(manufacturer, model)
|
|
model_hash = str(model_info.get("hash"))
|
|
callback = partial(self.download_profile, manufacturer, model, storage_path, model_hash)
|
|
await self.download_with_retry(callback)
|
|
self.profile_hashes[f"{manufacturer}/{model}"] = model_hash
|
|
await self.hass.async_add_executor_job(self._write_profile_hashes, self.profile_hashes)
|
|
except ProfileDownloadError as e:
|
|
path_exists, storage_path_exists = await self.hass.async_add_executor_job(
|
|
self._profile_paths_exist,
|
|
model_path,
|
|
storage_path,
|
|
)
|
|
if not path_exists:
|
|
if storage_path_exists:
|
|
await self.hass.async_add_executor_job(shutil.rmtree, storage_path) # pragma: no cover
|
|
raise e
|
|
_LOGGER.debug("Failed to download profile, falling back to local profile")
|
|
|
|
@staticmethod
|
|
def _profile_paths_exist(model_path: str, storage_path: str) -> tuple[bool, bool]:
|
|
"""Check profile paths from the executor."""
|
|
return os.path.exists(model_path), os.path.exists(storage_path)
|
|
|
|
async def _load_model_json(self, model_path: str) -> dict:
|
|
"""Load the JSON data from the model file."""
|
|
|
|
def _load_json() -> dict[str, Any]:
|
|
with open(model_path) as f:
|
|
return cast(dict[str, Any], json.load(f))
|
|
|
|
return await self.hass.async_add_executor_job(_load_json)
|
|
|
|
async def _handle_json_decode_error(
|
|
self,
|
|
error: JSONDecodeError,
|
|
manufacturer: str,
|
|
model: str,
|
|
retry_count: int,
|
|
) -> tuple[dict, str] | None:
|
|
"""Handle JSON decode errors with retry logic."""
|
|
_LOGGER.error("model.json file is not valid JSON for manufacturer: %s, model: %s", manufacturer, model)
|
|
if retry_count < 2:
|
|
_LOGGER.debug("Retrying to load model.json file")
|
|
return await self.load_model(manufacturer, model, True, retry_count + 1)
|
|
raise LibraryLoadingError("Failed to load model.json file") from error
|
|
|
|
def get_storage_path(self, manufacturer: str, model: str) -> str:
|
|
"""Retrieve the storage path for a given manufacturer and model."""
|
|
return str(self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, manufacturer, model))
|
|
|
|
async def download_with_retry(
|
|
self,
|
|
callback: Callable[[], Coroutine[Any, Any, None | dict[str, Any]]],
|
|
) -> None | dict[str, Any]:
|
|
"""Download a file from a remote endpoint with retries"""
|
|
max_retries = 3
|
|
retry_count = 0
|
|
|
|
while retry_count < max_retries:
|
|
try:
|
|
return await callback()
|
|
except (ClientError, TimeoutError, ProfileDownloadError) as e:
|
|
_LOGGER.debug(e)
|
|
retry_count += 1
|
|
if retry_count == max_retries:
|
|
raise ProfileDownloadError(
|
|
f"Failed to download even after {max_retries} retries, falling back to local copy",
|
|
) from e
|
|
|
|
await asyncio.sleep(self.retry_timeout)
|
|
_LOGGER.warning("Failed to download, retrying... (Attempt %d of %d)", retry_count + 1, max_retries)
|
|
return None # pragma: no cover
|
|
|
|
async def download_profile(self, manufacturer: str, model: str, storage_path: str, model_hash: str) -> None:
|
|
"""
|
|
Download the profile from Github using the Powercalc download API
|
|
Saves the profile to manufacturer/model directory in .storage/powercalc_profiles folder
|
|
"""
|
|
|
|
_LOGGER.debug("Downloading profile: %s/%s from github", manufacturer, model)
|
|
|
|
endpoint = f"{ENDPOINT_DOWNLOAD}/{manufacturer}/{model}"
|
|
|
|
def _save_file(data: bytes, directory: str) -> None:
|
|
"""Save file from Github to local storage directory"""
|
|
path = os.path.join(storage_path, directory)
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
with open(path, "wb") as f:
|
|
f.write(data)
|
|
|
|
session = async_get_clientsession(self.hass)
|
|
|
|
try:
|
|
async with asyncio.timeout(TIMEOUT_SECONDS):
|
|
async with session.get(endpoint, params={"hash": model_hash}) as resp:
|
|
if resp.status != 200:
|
|
raise ProfileDownloadError(f"Failed to download profile: {manufacturer}/{model}")
|
|
resources = await resp.json()
|
|
|
|
await self.hass.async_add_executor_job(lambda: os.makedirs(storage_path, exist_ok=True))
|
|
|
|
# Download the files
|
|
for resource in resources:
|
|
url = resource.get("url")
|
|
async with session.get(url) as resp:
|
|
if resp.status != 200:
|
|
raise ProfileDownloadError(f"Failed to download github URL: {url}")
|
|
|
|
contents = await resp.read()
|
|
await self.hass.async_add_executor_job(_save_file, contents, resource.get("path"))
|
|
except (TimeoutError, aiohttp.ClientError) as e:
|
|
raise ProfileDownloadError(f"Failed to download profile: {manufacturer}/{model}") from e
|
|
|
|
def _get_profile_hashes_path(self) -> str:
|
|
"""Retrieve the local storage path for the profile hashes file."""
|
|
return str(self.hass.config.path(STORAGE_DIR, BUILT_IN_LIBRARY_DIR, ".profile_hashes"))
|
|
|
|
def _load_profile_hashes(self) -> dict[str, str]:
|
|
"""Load profile hashes from local storage"""
|
|
|
|
path = self._get_profile_hashes_path()
|
|
if not os.path.exists(path):
|
|
return {}
|
|
|
|
with open(path) as f:
|
|
return json.load(f) # type: ignore[no-any-return]
|
|
|
|
def _write_profile_hashes(self, hashes: dict[str, str]) -> None:
|
|
"""Write profile hashes to local storage"""
|
|
|
|
path = self._get_profile_hashes_path()
|
|
with open(path, "w") as json_file:
|
|
json.dump(hashes, json_file, indent=4)
|