goodwe-addon/goodwe_controller/app/haentities.py

"""UX-13: Home Assistant's entities, for the settings page's entity picker.

A read of `GET /api/states` through the add-on's own session (hass.states()),
cached for 30 s and trimmed to what a picker shows. Nothing here writes, so
UX-02 §6 holds. When Home Assistant does not answer, the last good list is
served with its age, or an empty one that says why: the route never fails.
"""

import asyncio
import time

CACHE_S = 30.0
DOWN_TEXT = "Home Assistant is not answering, so the entity list is empty."
STALE_TEXT = "Home Assistant is not answering, so this list may be old."
_POWER_UNITS = ("mW", "W", "kW", "MW", "GW", "TW")
_TO_W = {"mW": 0.001, "W": 1.0, "kW": 1000.0, "MW": 1e6}
_ENERGY_UNITS = ("Wh", "kWh", "MWh", "GWh")


def row(state: dict) -> dict | None:
    """One state object -> {entity_id, name, state, unit, device_class}."""
    if not isinstance(state, dict) or not isinstance(state.get("entity_id"), str):
        return None
    attrs = state.get("attributes") if isinstance(state.get("attributes"), dict) else {}
    return {"entity_id": state["entity_id"],
            "name": str(attrs.get("friendly_name") or state["entity_id"]),
            "state": state.get("state"),
            "unit": attrs.get("unit_of_measurement"),
            "device_class": attrs.get("device_class")}


def watts(r: dict | None):
    """A power entity's reading in W, or None when it is not a number in a
    power unit (the setup wizard's kettle test)."""
    if not r or r.get("unit") not in _TO_W:
        return None
    try:
        value = float(r.get("state"))
    except (TypeError, ValueError):
        return None
    return value * _TO_W[r["unit"]] if value == value else None


def matches(r: dict, kind: str) -> bool:
    """Is `r` what a picker of this kind offers first? `kind` is one of
    uilayout.ENTITY_KINDS; "" or "any" offers everything."""
    unit, cls = r.get("unit"), r.get("device_class")
    if kind in ("power", "power_w"):
        return cls == "power" or unit in _POWER_UNITS
    if kind == "energy":
        return cls == "energy" or unit in _ENERGY_UNITS
    if kind == "pct":
        return unit == "%"
    if kind == "number":
        return r["entity_id"].startswith(("number.", "input_number."))
    return True


class EntityCache:
    """The last entity list read from Home Assistant, at most CACHE_S old."""

    def __init__(self, hass, clock=time.monotonic):
        self.hass = hass
        self.clock = clock
        self.rows: dict | None = None
        self.read_at: float | None = None
        self.tried_at: float | None = None
        self._lock = asyncio.Lock()

    async def get(self, max_age: float = CACHE_S
                  ) -> tuple[dict | None, float | None]:
        """(rows by entity id, age in s), or (None, None) if never read.
        `max_age` is how old a try may be before HA is asked again: the setup
        wizard's last step asks for a list at most a few seconds old."""
        async with self._lock:      # one HA read however many pickers ask
            now = self.clock()
            # A failed read also waits CACHE_S, so polls while HA is down do
            # not queue one 10 s read each behind this lock.
            if self.tried_at is None or now - self.tried_at >= max_age:
                self.tried_at = now
                states = await self.hass.states()
                if states is not None:
                    rows = (row(s) for s in states)
                    self.rows = {r["entity_id"]: r for r in rows if r}
                    self.read_at = now
            if self.read_at is None:
                return None, None
            return self.rows, now - self.read_at

    async def reading(self, entity_id: str) -> dict:
        """One entity read now, not from the list: the kettle test compares
        two readings seconds apart, which a 30 s list cannot. A read only."""
        r = row(await self.hass.state(entity_id)) if entity_id else None
        return {"entity_id": entity_id, "found": r is not None,
                "state": r and r["state"], "unit": r and r["unit"],
                "watts": watts(r)}

    def peek(self) -> dict | None:
        """The cached rows if they are fresh, without asking Home Assistant."""
        if self.read_at is None or self.clock() - self.read_at >= CACHE_S:
            return None
        return self.rows

    async def payload(self, kind: str = "") -> dict:
        """The route's JSON: matching entities first, then the rest."""
        rows, age = await self.get()
        if rows is None:
            return {"available": False, "text": DOWN_TEXT, "entities": []}
        ordered = sorted(rows.values(), key=lambda r: r["entity_id"])
        first = [dict(r, match=True) for r in ordered if matches(r, kind)]
        rest = [dict(r, match=False) for r in ordered if not matches(r, kind)]
        out = {"available": True, "age_s": round(age, 1), "entities": first + rest}
        if age >= CACHE_S:
            out["text"] = STALE_TEXT
        return out