from functools import partial import json import logging import os from typing import Any, cast from homeassistant.core import HomeAssistant from custom_components.powercalc.power_profile.error import LibraryLoadingError from custom_components.powercalc.power_profile.loader.protocol import Loader from custom_components.powercalc.power_profile.power_profile import DeviceType, DiscoveryBy, PowerProfile _LOGGER = logging.getLogger(__name__) class LocalLoader(Loader): def __init__(self, hass: HomeAssistant, directory: str, is_custom_directory: bool = False) -> None: self._is_custom_directory = is_custom_directory self._data_directory = directory self._hass = hass self._manufacturer_model_listing: dict[str, dict[str, PowerProfile]] = {} async def initialize(self) -> None: """Initialize the loader.""" if not self._is_custom_directory: await self._hass.async_add_executor_job(self._load_custom_library) async def get_manufacturer_listing( self, device_types: set[DeviceType] | None, discovery_by: DiscoveryBy | None = None, ) -> set[tuple[str, str]]: """Get listing of all available manufacturers or filtered by model device_type.""" if device_types is None: if discovery_by is None: return {(manufacturer, manufacturer) for manufacturer in self._manufacturer_model_listing} return { (manufacturer, manufacturer) for manufacturer, profiles in self._manufacturer_model_listing.items() if any(profile.discovery_by == discovery_by for profile in profiles.values()) } manufacturers: set[tuple[str, str]] = set() for manufacturer in self._manufacturer_model_listing: models = await self.get_model_listing(manufacturer, device_types, discovery_by) if not models: continue manufacturers.add((manufacturer, manufacturer)) return manufacturers async def find_manufacturers(self, search: str) -> set[str]: """Check if a manufacturer is available.""" _search = search.lower() manufacturer_list = self._manufacturer_model_listing.keys() if _search in manufacturer_list: return {_search} return set() 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 for a given manufacturer. param manufacturer: manufacturer always handled in lower case param device_type: models of the manufacturer will be filtered by DeviceType, models without assigned device_type will be handled as DeviceType.LIGHT. None will return all models of a manufacturer. returns: Set[tuple[str, str]] of (model_id, model_name) """ found_models: set[tuple[str, str]] = set() models = self._manufacturer_model_listing.get(manufacturer.lower()) if not models: return found_models for profile in models.values(): if device_types and profile.device_type not in device_types: continue if discovery_by and profile.discovery_by != discovery_by: continue found_models.add((profile.model, profile.name or profile.model)) return found_models async def load_model(self, manufacturer: str, model: str) -> tuple[dict, str] | None: """Load a model.json file from disk for a given manufacturer.lower() and model.lower() by querying the custom library. If self._is_custom_directory == true model.json will be loaded directly from there. returns: tuple[dict, str] model.json as dictionary and model as lower case returns: None when manufacturer, model or model path not found raises LibraryLoadingError: model.json not found """ _manufacturer = manufacturer.lower() _model = model.lower() if self._is_custom_directory: model_path, model_json = await self._hass.async_add_executor_job( partial(self._load_custom_model, _manufacturer, _model), ) return model_json, model_path lib_models = self._manufacturer_model_listing.get(_manufacturer) if lib_models is None: return None lib_model = lib_models.get(_model) if lib_model is None: return None model_path = lib_model.get_model_directory() model_json = lib_model.json_data return model_json, model_path async def find_model(self, manufacturer: str, search: set[str]) -> list[str]: """Find a model for a given manufacturer. Also must check aliases.""" _manufacturer = manufacturer.lower() models = self._manufacturer_model_listing.get(_manufacturer) if not models: return [] search_lower = {phrase.lower() for phrase in search} profile = next((models[model] for model in models if model.lower() in search_lower), None) return [profile.model] if profile else [] async def find_model_migration(self, manufacturer: str, model: str) -> str | None: """Local custom libraries do not support metadata-driven legacy profile migrations.""" return None def _load_custom_library(self) -> None: """Loading custom models and aliases from file system. Manufacturer directories without model directories and model.json files within are not loaded. Same is with model directories without model.json files. """ base_path = self._data_directory if not os.path.exists(base_path): _LOGGER.error("Custom library directory does not exist: %s", base_path) return self._manufacturer_model_listing.clear() for manufacturer_dir in next(os.walk(base_path))[1]: manufacturer_path = os.path.join(base_path, manufacturer_dir) manufacturer = manufacturer_dir.lower() for model_dir in next(os.walk(manufacturer_path))[1]: if model_dir.startswith("."): continue model_path = os.path.join(manufacturer_path, model_dir) model_json_path = os.path.join(model_path, "model.json") if not os.path.exists(model_json_path): _LOGGER.warning("model.json should exist in %s!", model_path) continue model_json = self._load_json(model_json_path) profile = PowerProfile( self._hass, manufacturer=manufacturer, model=model_dir, directory=model_path, json_data=model_json, ) self._add_profile_to_library(profile) for alias in profile.aliases: self._add_profile_to_library( PowerProfile( self._hass, manufacturer=manufacturer, model=alias, directory=model_path, json_data=model_json, ), ) def _add_profile_to_library(self, profile: PowerProfile) -> None: """Add profile to the library lookup dictionary.""" manufacturer = profile.manufacturer if self._manufacturer_model_listing.get(manufacturer) is None: self._manufacturer_model_listing[manufacturer] = {} search_key = profile.model.lower() if self._manufacturer_model_listing[manufacturer].get(search_key): _LOGGER.error( "Double entry manufacturer/model in custom library: %s/%s", profile.manufacturer, profile.model, ) return self._manufacturer_model_listing[manufacturer].update({search_key: profile}) def _load_custom_model(self, manufacturer: str, model: str) -> tuple[str, dict[str, Any]]: """Load model.json from a directly configured custom model directory.""" model_path = os.path.join(self._data_directory) model_json_path = os.path.join(model_path, "model.json") if not os.path.exists(model_json_path): raise LibraryLoadingError( f"model.json not found for manufacturer {manufacturer} and model {model} in path {model_json_path}", ) return model_path, self._load_json(model_json_path) def _load_json(self, model_json_path: str) -> dict[str, Any]: """Load model.json file for a given model.""" with open(model_json_path) as file: return cast(dict[str, Any], json.load(file))