395 lines
16 KiB
Python
395 lines
16 KiB
Python
from copy import deepcopy
|
|
import json
|
|
import os
|
|
import re
|
|
from typing import Any, NamedTuple, cast
|
|
|
|
from homeassistant.core import HomeAssistant
|
|
from homeassistant.helpers.singleton import singleton
|
|
|
|
from custom_components.powercalc.common import SourceEntity
|
|
from custom_components.powercalc.const import CONF_DISABLE_LIBRARY_DOWNLOAD, DOMAIN, DOMAIN_CONFIG
|
|
from custom_components.powercalc.helpers import (
|
|
build_related_entity_placeholder_not_found_message,
|
|
collect_placeholders,
|
|
iter_related_entity_placeholders,
|
|
replace_placeholders,
|
|
resolve_related_entity_placeholder,
|
|
)
|
|
|
|
from .error import LibraryError
|
|
from .loader.composite import CompositeLoader
|
|
from .loader.local import LocalLoader
|
|
from .loader.protocol import Loader, ModelMetadata
|
|
from .loader.remote import RemoteLoader
|
|
from .power_profile import DeviceType, DiscoveryBy, PowerProfile
|
|
|
|
LEGACY_CUSTOM_DATA_DIRECTORY = "powercalc-custom-models"
|
|
CUSTOM_DATA_DIRECTORY = "powercalc/profiles"
|
|
|
|
|
|
def load_sub_profile_data(base_dir: str) -> list[tuple[str, dict[str, Any]]]:
|
|
"""Load sub-profile JSON blobs from disk."""
|
|
sub_dirs = next(os.walk(base_dir))[1]
|
|
result = []
|
|
for sub_dir in sub_dirs:
|
|
json_path = os.path.join(base_dir, sub_dir, "model.json")
|
|
if os.path.isfile(json_path):
|
|
with open(json_path, encoding="utf-8") as f:
|
|
json_data = cast(dict[str, Any], json.load(f))
|
|
else:
|
|
json_data = {}
|
|
result.append((sub_dir, json_data))
|
|
return sorted(result, key=lambda item: item[0])
|
|
|
|
|
|
class ProfileLibrary:
|
|
def __init__(self, hass: HomeAssistant, loader: Loader) -> None:
|
|
self._hass = hass
|
|
self._loader = loader
|
|
self._profiles: dict[str, list[PowerProfile]] = {}
|
|
self._manufacturer_models: dict[str, set[tuple[str, str]]] = {}
|
|
self._sub_profile_data: dict[str, list[tuple[str, dict[str, Any]]]] = {}
|
|
self._found_models: dict[ModelInfo, list[ModelInfo]] = {}
|
|
|
|
async def initialize(self, prefer_cached: bool = False) -> None:
|
|
"""Initialize the underlying loaders, see `Loader.initialize` for `prefer_cached`."""
|
|
self._sub_profile_data.clear()
|
|
self._found_models.clear()
|
|
await self._loader.initialize(prefer_cached)
|
|
|
|
@property
|
|
def discovery_low_priority_domains(self) -> set[str]:
|
|
"""Get integration domains that are the least preferred source for discovery.
|
|
|
|
Devices behind these integrations are only discovered when no other integration
|
|
represents them, and their entities are never discovered by entity discovery.
|
|
"""
|
|
return self._loader.get_discovery_low_priority_domains()
|
|
|
|
@staticmethod
|
|
@singleton("powercalc_library")
|
|
async def factory(hass: HomeAssistant) -> ProfileLibrary:
|
|
"""
|
|
Creates and loads the profile library.
|
|
Make sure we have a single instance throughout the application.
|
|
"""
|
|
library = ProfileLibrary(hass, ProfileLibrary.create_loader(hass))
|
|
# Startup must not block on the download API. The periodic library update refreshes it.
|
|
await library.initialize(prefer_cached=True)
|
|
return library
|
|
|
|
@staticmethod
|
|
def create_loader(hass: HomeAssistant, skip_remote_loader: bool = False) -> Loader:
|
|
loaders: list[Loader] = [
|
|
LocalLoader(hass, data_dir)
|
|
for data_dir in [
|
|
os.path.join(hass.config.config_dir, LEGACY_CUSTOM_DATA_DIRECTORY),
|
|
os.path.join(hass.config.config_dir, CUSTOM_DATA_DIRECTORY),
|
|
os.path.join(os.path.dirname(__file__), "../custom_data"),
|
|
]
|
|
if os.path.exists(data_dir)
|
|
]
|
|
|
|
domain_config = hass.data.get(DOMAIN, {})
|
|
global_config = domain_config.get(DOMAIN_CONFIG, {})
|
|
disable_library_download: bool = bool(global_config.get(CONF_DISABLE_LIBRARY_DOWNLOAD, False))
|
|
if not disable_library_download and not skip_remote_loader:
|
|
loaders.append(RemoteLoader(hass))
|
|
|
|
return CompositeLoader(loaders)
|
|
|
|
async def get_manufacturer_listing(
|
|
self,
|
|
device_types: set[DeviceType] | None = None,
|
|
discovery_by: DiscoveryBy | None = None,
|
|
) -> list[tuple[str, str]]:
|
|
"""Get listing of available manufacturers."""
|
|
manufacturers = await self._loader.get_manufacturer_listing(device_types, discovery_by)
|
|
return sorted(manufacturers)
|
|
|
|
async def get_model_listing(
|
|
self,
|
|
manufacturer: str,
|
|
device_types: set[DeviceType] | None = None,
|
|
discovery_by: DiscoveryBy | None = None,
|
|
) -> list[tuple[str, str]]:
|
|
"""Get listing of available models and display names for a given manufacturer."""
|
|
|
|
resolved_manufacturers = await self._loader.find_manufacturers(manufacturer)
|
|
if not resolved_manufacturers:
|
|
return []
|
|
|
|
all_models: list[tuple[str, str]] = []
|
|
for manufacturer in resolved_manufacturers:
|
|
cache_key = f"{manufacturer}/{device_types}/{discovery_by}"
|
|
cached_models = self._manufacturer_models.get(cache_key)
|
|
if cached_models:
|
|
all_models.extend(sorted(cached_models))
|
|
continue
|
|
models = await self._loader.get_model_listing(manufacturer, device_types, discovery_by)
|
|
self._manufacturer_models[cache_key] = models
|
|
all_models.extend(sorted(models))
|
|
|
|
return sorted(all_models, key=lambda model: model[0])
|
|
|
|
async def get_profile(
|
|
self,
|
|
model_info: ModelInfo,
|
|
source_entity: SourceEntity | None = None,
|
|
custom_directory: str | None = None,
|
|
variables: dict[str, str] | None = None,
|
|
process_variables: bool = True,
|
|
model_resolved: bool = False,
|
|
) -> PowerProfile:
|
|
"""Get a power profile for a given manufacturer and model.
|
|
|
|
Pass `model_resolved` when `model_info` already comes out of `find_models`, to skip
|
|
looking the model up in the library a second time.
|
|
"""
|
|
# Support multiple LUT in subdirectories
|
|
sub_profile = None
|
|
if "/" in model_info.model:
|
|
(model, sub_profile) = model_info.model.split("/", 1)
|
|
model_info = ModelInfo(model_info.manufacturer, model, model_info.model_id)
|
|
|
|
if not custom_directory and not model_resolved:
|
|
models = await self.find_models(model_info)
|
|
if not models:
|
|
raise LibraryError(f"Model {model_info.manufacturer} {model_info.model} not found")
|
|
model_info = next(iter(models))
|
|
|
|
profile = await self.create_power_profile(
|
|
model_info,
|
|
source_entity,
|
|
custom_directory,
|
|
variables,
|
|
process_variables,
|
|
)
|
|
|
|
if sub_profile:
|
|
await profile.select_sub_profile(sub_profile)
|
|
|
|
return profile
|
|
|
|
async def create_power_profile(
|
|
self,
|
|
model_info: ModelInfo,
|
|
source_entity: SourceEntity | None = None,
|
|
custom_directory: str | None = None,
|
|
variables: dict[str, str] | None = None,
|
|
process_variables: bool = True,
|
|
) -> PowerProfile:
|
|
"""Create a power profile object from the model JSON data."""
|
|
|
|
json_data, directory = await self._load_model_data(model_info.manufacturer, model_info.model, custom_directory)
|
|
json_data = self._process_profile_json(json_data, variables or {}, source_entity, process_variables)
|
|
|
|
if linked_profile := json_data.get("linked_profile", json_data.get("linked_lut")):
|
|
linked_manufacturer, linked_model = linked_profile.split("/")
|
|
linked_json_data, directory = await self._load_model_data(
|
|
linked_manufacturer,
|
|
linked_model,
|
|
custom_directory,
|
|
)
|
|
json_data.update(linked_json_data)
|
|
|
|
raw_sub_profiles = self._sub_profile_data.get(directory)
|
|
if raw_sub_profiles is None:
|
|
raw_sub_profiles = await self._hass.async_add_executor_job(load_sub_profile_data, directory)
|
|
self._sub_profile_data[directory] = raw_sub_profiles
|
|
|
|
sub_profiles = [
|
|
(
|
|
sub_dir,
|
|
self._process_profile_json(sub_profile_json, variables or {}, source_entity, process_variables),
|
|
)
|
|
for sub_dir, sub_profile_json in raw_sub_profiles
|
|
]
|
|
|
|
return await self._create_power_profile_instance(
|
|
model_info.manufacturer,
|
|
model_info.model,
|
|
directory,
|
|
json_data,
|
|
sub_profiles,
|
|
)
|
|
|
|
def _process_profile_json(
|
|
self,
|
|
json_data: dict[str, Any],
|
|
variables: dict[str, str],
|
|
source_entity: SourceEntity | None,
|
|
process_variables: bool,
|
|
) -> dict[str, Any]:
|
|
if not process_variables:
|
|
# json_data is retrieved from cache, so we need to copy it to avoid modifying the cache
|
|
return json_data.copy()
|
|
|
|
# replace_placeholders rewrites nested dicts and lists in place, so a shallow copy is not
|
|
# enough here. Without a deep copy the substituted values leak into the cached profile data
|
|
# and the next profile built from the same model would reuse them.
|
|
json_data = deepcopy(json_data)
|
|
|
|
if json_data.get("fields"): # When custom fields in profile are defined, make sure all variables are passed
|
|
self.validate_variables(json_data, variables)
|
|
|
|
placeholders = collect_placeholders(json_data)
|
|
replacements = self.compute_replacement_variables(placeholders, variables.copy(), source_entity)
|
|
return cast(dict[str, Any], replace_placeholders(json_data, replacements))
|
|
|
|
def compute_replacement_variables(
|
|
self,
|
|
placeholders: set[str],
|
|
variables: dict[str, str],
|
|
source_entity: SourceEntity | None,
|
|
) -> dict[str, str]:
|
|
variables = variables or {}
|
|
|
|
if source_entity:
|
|
if "entity" in placeholders:
|
|
variables["entity"] = source_entity.entity_id
|
|
|
|
for placeholder in iter_related_entity_placeholders(placeholders):
|
|
related_entity = resolve_related_entity_placeholder(
|
|
self._hass,
|
|
placeholder,
|
|
source_entity=source_entity,
|
|
)
|
|
if not related_entity:
|
|
raise LibraryError(
|
|
build_related_entity_placeholder_not_found_message(placeholder, source_entity.entity_id),
|
|
)
|
|
variables[placeholder] = related_entity
|
|
|
|
return variables
|
|
|
|
@staticmethod
|
|
def validate_variables(json_data: dict[str, Any], variables: dict[str, str]) -> None:
|
|
fields = json_data.get("fields", {}).keys()
|
|
|
|
# Check if all variables are valid for the model
|
|
for variable in variables:
|
|
if variable not in fields and variable != "entity":
|
|
raise LibraryError(f"Variable {variable} is not valid for this model")
|
|
|
|
# Check if all fields have corresponding variables
|
|
missing_fields = [field for field in fields if field not in variables]
|
|
if missing_fields:
|
|
raise LibraryError(f"Missing variables for fields: {', '.join(missing_fields)}")
|
|
|
|
async def find_manufacturers(self, manufacturer: str) -> set[str]:
|
|
"""Resolve the manufacturer, either from the model info or by loading it."""
|
|
return await self._loader.find_manufacturers(manufacturer)
|
|
|
|
async def get_model_metadata(self, model_info: ModelInfo) -> ModelMetadata | None:
|
|
"""Return discovery metadata for an already resolved model, without building the profile."""
|
|
return await self._loader.get_model_metadata(model_info.manufacturer, model_info.model)
|
|
|
|
async def find_models(self, model_info: ModelInfo) -> list[ModelInfo]:
|
|
"""Resolve the model identifier, searching for it if no custom directory is provided.
|
|
|
|
Discovery resolves the same handful of models for every entity of a device, so the
|
|
result is memoized until the library is reloaded.
|
|
"""
|
|
if model_info in self._found_models:
|
|
return self._found_models[model_info]
|
|
|
|
found = await self._find_models(model_info)
|
|
self._found_models[model_info] = found
|
|
return found
|
|
|
|
async def _find_models(self, model_info: ModelInfo) -> list[ModelInfo]:
|
|
"""Search the loaders for all models matching the given model info."""
|
|
search: set[str] = set()
|
|
for model_identifier in (model_info.model_id, model_info.model):
|
|
if model_identifier:
|
|
model_identifier = model_identifier.replace("#slash#", "/")
|
|
search.update(
|
|
{
|
|
model_identifier,
|
|
model_identifier.lower(),
|
|
re.sub(r"^(.*)\(([^()]+)\)$", r"\2", model_identifier),
|
|
},
|
|
)
|
|
if "/" in model_identifier:
|
|
search.update(model_identifier.split("/"))
|
|
|
|
manufacturers = await self._loader.find_manufacturers(model_info.manufacturer)
|
|
if not manufacturers:
|
|
return []
|
|
|
|
found_models: list[ModelInfo] = []
|
|
for manufacturer in manufacturers:
|
|
models = await self._loader.find_model(manufacturer, search)
|
|
if models:
|
|
found_models.extend(ModelInfo(manufacturer, model) for model in models)
|
|
|
|
return list(dict.fromkeys(found_models))
|
|
|
|
async def find_model_migration(self, model_info: ModelInfo) -> ModelInfo | None:
|
|
"""Resolve a legacy canonical model id to its replacement using library metadata."""
|
|
manufacturers = await self._loader.find_manufacturers(model_info.manufacturer)
|
|
if not manufacturers:
|
|
return None
|
|
|
|
matches: set[ModelInfo] = set()
|
|
for manufacturer in manufacturers:
|
|
migrated_model = await self._loader.find_model_migration(manufacturer, model_info.model)
|
|
if migrated_model:
|
|
matches.add(ModelInfo(manufacturer, migrated_model))
|
|
|
|
if len(matches) != 1:
|
|
return None
|
|
|
|
return next(iter(matches))
|
|
|
|
async def _load_model_data(
|
|
self,
|
|
manufacturer: str,
|
|
model: str,
|
|
custom_directory: str | None,
|
|
) -> tuple[dict[str, Any], str]:
|
|
"""Load the model data from the appropriate directory."""
|
|
loader = (
|
|
LocalLoader(self._hass, custom_directory, is_custom_directory=True) if custom_directory else self._loader
|
|
)
|
|
result = await loader.load_model(manufacturer, model)
|
|
if not result:
|
|
raise LibraryError(f"Model {manufacturer} {model} not found")
|
|
|
|
return result
|
|
|
|
async def _create_power_profile_instance(
|
|
self,
|
|
manufacturer: str,
|
|
model: str,
|
|
directory: str,
|
|
json_data: dict[str, Any],
|
|
sub_profiles: list[tuple[str, dict[str, Any]]] | None = None,
|
|
) -> PowerProfile:
|
|
"""Create and initialize the PowerProfile object."""
|
|
profile = PowerProfile(
|
|
self._hass,
|
|
manufacturer=manufacturer,
|
|
model=model,
|
|
directory=directory,
|
|
json_data=json_data,
|
|
sub_profiles=sub_profiles,
|
|
)
|
|
|
|
if not profile.sub_profile and profile.sub_profile_select:
|
|
await profile.select_sub_profile(profile.sub_profile_select.default)
|
|
|
|
return profile
|
|
|
|
def get_loader(self) -> Loader:
|
|
return self._loader
|
|
|
|
|
|
class ModelInfo(NamedTuple):
|
|
manufacturer: str
|
|
model: str
|
|
# Starting from HA 2024.8 we can use model_id to identify the model
|
|
model_id: str | None = None
|