data.py

Go to the documentation of this file.

"""Energy data."""
 
from __future__ import annotations
 
import asyncio
from collections import Counter
from collections.abc import Awaitable, Callable
from typing import Literal, TypedDict
 
import voluptuous as vol
 
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import config_validation as cv, singleton, storage
 
from .const import DOMAIN
 
STORAGE_VERSION = 1
STORAGE_KEY = DOMAIN
 
 
@singleton.singleton(f"{DOMAIN}_manager")
async def async_get_manager(hass: HomeAssistant) -> EnergyManager:
    """Return an initialized data manager."""
    manager = EnergyManager(hass)
    await manager.async_initialize()
    return manager
 
 
class FlowFromGridSourceType(TypedDict):
    """Dictionary describing the 'from' stat for the grid source."""
 
    # statistic_id of a an energy meter (kWh)
    stat_energy_from: str
 
    # statistic_id of costs ($) incurred from the energy meter
    # If set to None and entity_energy_price or number_energy_price are configured,
    # an EnergyCostSensor will be automatically created
    stat_cost: str | None
 
    # Used to generate costs if stat_cost is set to None
    entity_energy_price: str | None  # entity_id of an entity providing price ($/kWh)
    number_energy_price: float | None  # Price for energy ($/kWh)
 
 
class FlowToGridSourceType(TypedDict):
    """Dictionary describing the 'to' stat for the grid source."""
 
    # kWh meter
    stat_energy_to: str
 
    # statistic_id of compensation ($) received for contributing back
    # If set to None and entity_energy_price or number_energy_price are configured,
    # an EnergyCostSensor will be automatically created
    stat_compensation: str | None
 
    # Used to generate costs if stat_compensation is set to None
    entity_energy_price: str | None  # entity_id of an entity providing price ($/kWh)
    number_energy_price: float | None  # Price for energy ($/kWh)
 
 
class GridSourceType(TypedDict):
    """Dictionary holding the source of grid energy consumption."""
 
    type: Literal["grid"]
 
    flow_from: list[FlowFromGridSourceType]
    flow_to: list[FlowToGridSourceType]
 
    cost_adjustment_day: float
 
 
class SolarSourceType(TypedDict):
    """Dictionary holding the source of energy production."""
 
    type: Literal["solar"]
 
    stat_energy_from: str
    config_entry_solar_forecast: list[str] | None
 
 
class BatterySourceType(TypedDict):
    """Dictionary holding the source of battery storage."""
 
    type: Literal["battery"]
 
    stat_energy_from: str
    stat_energy_to: str
 
 
class GasSourceType(TypedDict):
    """Dictionary holding the source of gas consumption."""
 
    type: Literal["gas"]
 
    stat_energy_from: str
 
    # statistic_id of costs ($) incurred from the gas meter
    # If set to None and entity_energy_price or number_energy_price are configured,
    # an EnergyCostSensor will be automatically created
    stat_cost: str | None
 
    # Used to generate costs if stat_cost is set to None
    entity_energy_price: str | None  # entity_id of an entity providing price ($/m³)
    number_energy_price: float | None  # Price for energy ($/m³)
 
 
class WaterSourceType(TypedDict):
    """Dictionary holding the source of water consumption."""
 
    type: Literal["water"]
 
    stat_energy_from: str
 
    # statistic_id of costs ($) incurred from the water meter
    # If set to None and entity_energy_price or number_energy_price are configured,
    # an EnergyCostSensor will be automatically created
    stat_cost: str | None
 
    # Used to generate costs if stat_cost is set to None
    entity_energy_price: str | None  # entity_id of an entity providing price ($/m³)
    number_energy_price: float | None  # Price for energy ($/m³)
 
 
type SourceType = (
    GridSourceType
    | SolarSourceType
    | BatterySourceType
    | GasSourceType
    | WaterSourceType
)
 
 
class DeviceConsumption(TypedDict):
    """Dictionary holding the source of individual device consumption."""
 
    # This is an ever increasing value
    stat_consumption: str
 
    # An optional custom name for display in energy graphs
    name: str | None
 
 
class EnergyPreferences(TypedDict):
    """Dictionary holding the energy data."""
 
    energy_sources: list[SourceType]
    device_consumption: list[DeviceConsumption]
 
 
class EnergyPreferencesUpdate(EnergyPreferences, total=False):
    """all types optional."""
 
 
def _flow_from_ensure_single_price(
    val: FlowFromGridSourceType,
) -> FlowFromGridSourceType:
    """Ensure we use a single price source."""
    if (
        val["entity_energy_price"] is not None
        and val["number_energy_price"] is not None
    ):
        raise vol.Invalid("Define either an entity or a fixed number for the price")
 
    return val
 
 
FLOW_FROM_GRID_SOURCE_SCHEMA = vol.All(
    vol.Schema(
        {
            vol.Required("stat_energy_from"): str,
            vol.Optional("stat_cost"): vol.Any(str, None),
            # entity_energy_from was removed in HA Core 2022.10
            vol.Remove("entity_energy_from"): vol.Any(str, None),
            vol.Optional("entity_energy_price"): vol.Any(str, None),
            vol.Optional("number_energy_price"): vol.Any(vol.Coerce(float), None),
        }
    ),
    _flow_from_ensure_single_price,
)
 
 
FLOW_TO_GRID_SOURCE_SCHEMA = vol.Schema(
    {
        vol.Required("stat_energy_to"): str,
        vol.Optional("stat_compensation"): vol.Any(str, None),
        # entity_energy_to was removed in HA Core 2022.10
        vol.Remove("entity_energy_to"): vol.Any(str, None),
        vol.Optional("entity_energy_price"): vol.Any(str, None),
        vol.Optional("number_energy_price"): vol.Any(vol.Coerce(float), None),
    }
)
 
 
def _generate_unique_value_validator(key: str) -> Callable[[list[dict]], list[dict]]:
    """Generate a validator that ensures a value is only used once."""
 
    def validate_uniqueness(
        val: list[dict],
    ) -> list[dict]:
        """Ensure that the user doesn't add duplicate values."""
        counts = Counter(flow_from[key] for flow_from in val)
 
        for value, count in counts.items():
            if count > 1:
                raise vol.Invalid(f"Cannot specify {value} more than once")
 
        return val
 
    return validate_uniqueness
 
 
GRID_SOURCE_SCHEMA = vol.Schema(
    {
        vol.Required("type"): "grid",
        vol.Required("flow_from"): vol.All(
            [FLOW_FROM_GRID_SOURCE_SCHEMA],
            _generate_unique_value_validator("stat_energy_from"),
        ),
        vol.Required("flow_to"): vol.All(
            [FLOW_TO_GRID_SOURCE_SCHEMA],
            _generate_unique_value_validator("stat_energy_to"),
        ),
        vol.Required("cost_adjustment_day"): vol.Coerce(float),
    }
)
SOLAR_SOURCE_SCHEMA = vol.Schema(
    {
        vol.Required("type"): "solar",
        vol.Required("stat_energy_from"): str,
        vol.Optional("config_entry_solar_forecast"): vol.Any([str], None),
    }
)
BATTERY_SOURCE_SCHEMA = vol.Schema(
    {
        vol.Required("type"): "battery",
        vol.Required("stat_energy_from"): str,
        vol.Required("stat_energy_to"): str,
    }
)
GAS_SOURCE_SCHEMA = vol.Schema(
    {
        vol.Required("type"): "gas",
        vol.Required("stat_energy_from"): str,
        vol.Optional("stat_cost"): vol.Any(str, None),
        # entity_energy_from was removed in HA Core 2022.10
        vol.Remove("entity_energy_from"): vol.Any(str, None),
        vol.Optional("entity_energy_price"): vol.Any(str, None),
        vol.Optional("number_energy_price"): vol.Any(vol.Coerce(float), None),
    }
)
WATER_SOURCE_SCHEMA = vol.Schema(
    {
        vol.Required("type"): "water",
        vol.Required("stat_energy_from"): str,
        vol.Optional("stat_cost"): vol.Any(str, None),
        vol.Optional("entity_energy_price"): vol.Any(str, None),
        vol.Optional("number_energy_price"): vol.Any(vol.Coerce(float), None),
    }
)
 
 
def check_type_limits(value: list[SourceType]) -> list[SourceType]:
    """Validate that we don't have too many of certain types."""
    types = Counter([val["type"] for val in value])
 
    if types.get("grid", 0) > 1:
        raise vol.Invalid("You cannot have more than 1 grid source")
 
    return value
 
 
ENERGY_SOURCE_SCHEMA = vol.All(
    vol.Schema(
        [
            cv.key_value_schemas(
                "type",
                {
                    "grid": GRID_SOURCE_SCHEMA,
                    "solar": SOLAR_SOURCE_SCHEMA,
                    "battery": BATTERY_SOURCE_SCHEMA,
                    "gas": GAS_SOURCE_SCHEMA,
                    "water": WATER_SOURCE_SCHEMA,
                },
            )
        ]
    ),
    check_type_limits,
)
 
DEVICE_CONSUMPTION_SCHEMA = vol.Schema(
    {
        vol.Required("stat_consumption"): str,
        vol.Optional("name"): str,
    }
)
 
 
class EnergyManager:
    """Manage the instance energy prefs."""
 
    def __init__(self, hass: HomeAssistant) -> None:
        """Initialize energy manager."""
        self._hass_hass = hass
        self._store_store = storage.Store[EnergyPreferences](
            hass, STORAGE_VERSION, STORAGE_KEY
        )
        self.datadata: EnergyPreferences | None = None
        self._update_listeners: list[Callable[[], Awaitable]] = []
 
    async def async_initialize(self) -> None:
        """Initialize the energy integration."""
        self.datadata = await self._store_store.async_load()
 
    @staticmethod
    def default_preferences() -> EnergyPreferences:
        """Return default preferences."""
        return {
            "energy_sources": [],
            "device_consumption": [],
        }
 
    async def async_update(self, update: EnergyPreferencesUpdate) -> None:
        """Update the preferences."""
        if self.datadata is None:
            data = EnergyManager.default_preferences()
        else:
            data = self.datadata.copy()
 
        for key in (
            "energy_sources",
            "device_consumption",
        ):
            if key in update:
                data[key] = update[key]
 
        self.datadata = data
        self._store_store.async_delay_save(lambda: data, 60)
 
        if not self._update_listeners:
            return
 
        await asyncio.gather(*(listener() for listener in self._update_listeners))
 
    @callback
    def async_listen_updates(self, update_listener: Callable[[], Awaitable]) -> None:
        """Listen for data updates."""
        self._update_listeners.append(update_listener)