Skip to content

Commit

Permalink
Move to runtime_data
Browse files Browse the repository at this point in the history
  • Loading branch information
@astrandb
astrandb committed Dec 6, 2024
1 parent d72de97 commit 5977bc6
Show file tree
Hide file tree
Showing 4 changed files with 118 additions and 105 deletions.
71 changes: 44 additions & 27 deletions custom_components/weatherlink/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from __future__ import annotations

import asyncio
from dataclasses import dataclass
from datetime import timedelta
from email.utils import mktime_tz, parsedate_tz
import logging
Expand All @@ -27,6 +28,19 @@
)
from .pyweatherlink import WLHub, WLHubV2

type WLConfigEntry = ConfigEntry[WLData]


@dataclass
class WLData:
api: WLHub | WLHubV2
primary_tx_id: int
station_data: dict
sensors_metadata: dict
coordinator: DataUpdateCoordinator
current: dict


PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR]
SENSOR_TYPE_VUE_AND_VANTAGE_PRO = (
23,
Expand Down Expand Up @@ -65,42 +79,45 @@
_LOGGER = logging.getLogger(__name__)


async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async def async_setup_entry(hass: HomeAssistant, entry: WLConfigEntry) -> bool:
"""Set up Weatherlink from a config entry."""

hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][entry.entry_id] = {}
entry.runtime_data = WLData(
api=None,
primary_tx_id=1,
station_data={},
sensors_metadata={},
coordinator=None,
current={},
)

if entry.data[CONF_API_VERSION] == ApiVersion.API_V1:
hass.data[DOMAIN][entry.entry_id]["api"] = WLHub(
entry.runtime_data.api = WLHub(
websession=async_get_clientsession(hass),
username=entry.data[CONF_USERNAME],
password=entry.data[CONF_PASSWORD],
apitoken=entry.data[CONF_API_TOKEN],
)
hass.data[DOMAIN][entry.entry_id]["primary_tx_id"] = 1
entry.runtime_data.primary_tx_id = 1
tx_ids = [1]

if entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
hass.data[DOMAIN][entry.entry_id]["api"] = WLHubV2(
entry.runtime_data.api = WLHubV2(
websession=async_get_clientsession(hass),
station_id=entry.data[CONF_STATION_ID],
api_key_v2=entry.data[CONF_API_KEY_V2],
api_secret=entry.data[CONF_API_SECRET],
)
hass.data[DOMAIN][entry.entry_id]["station_data"] = await hass.data[DOMAIN][
entry.entry_id
]["api"].get_station()
entry.runtime_data.station_data = await entry.runtime_data.api.get_station()

all_sensors = await hass.data[DOMAIN][entry.entry_id]["api"].get_all_sensors()
all_sensors = await entry.runtime_data.api.get_all_sensors()

sensors = []
tx_ids = []
for sensor in all_sensors["sensors"]:
if (
sensor["station_id"]
== hass.data[DOMAIN][entry.entry_id]["station_data"]["stations"][0][
"station_id"
]
== entry.runtime_data.station_data["stations"][0]["station_id"]
):
sensors.append(sensor)
if (
Expand All @@ -109,11 +126,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
and sensor["tx_id"] not in tx_ids
):
tx_ids.append(sensor["tx_id"])
hass.data[DOMAIN][entry.entry_id]["sensors_metadata"] = sensors
entry.runtime_data.sensors_metadata = sensors
# todo Make primary_tx_id configurable by user - perhaps in config flow.
if len(tx_ids) == 0:
tx_ids = [1]
hass.data[DOMAIN][entry.entry_id]["primary_tx_id"] = min(tx_ids)
entry.runtime_data.primary_tx_id = min(tx_ids)
_LOGGER.debug("Primary tx_ids: %s", tx_ids)
coordinator = await get_coordinator(hass, entry)
if not coordinator.last_update_success:
Expand All @@ -125,10 +142,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return True


async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async def async_unload_entry(hass: HomeAssistant, entry: WLConfigEntry) -> bool:
"""Unload a config entry."""
if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
hass.data[DOMAIN].pop(entry.entry_id)
unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

return unload_ok

Expand All @@ -138,11 +154,12 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:

async def get_coordinator( # noqa: C901
hass: HomeAssistant,
entry: ConfigEntry,
entry: WLConfigEntry,
) -> DataUpdateCoordinator:
"""Get the data update coordinator."""
if "coordinator" in hass.data[DOMAIN][entry.entry_id]:
return hass.data[DOMAIN][entry.entry_id]["coordinator"]

if entry.runtime_data.coordinator is not None:
return entry.runtime_data.coordinator

def _preprocess(indata: str): # noqa: C901
outdata = {}
Expand Down Expand Up @@ -184,7 +201,7 @@ def _preprocess(indata: str): # noqa: C901
)

if entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
primary_tx_id = tx_id = hass.data[DOMAIN][entry.entry_id]["primary_tx_id"]
primary_tx_id = tx_id = entry.runtime_data.primary_tx_id
outdata.setdefault(tx_id, {})
outdata[DataKey.UUID] = indata["station_id_uuid"]
for sensor in indata["sensors"]:
Expand Down Expand Up @@ -523,26 +540,26 @@ def _preprocess(indata: str): # noqa: C901
return outdata

async def async_fetch():
api = hass.data[DOMAIN][entry.entry_id]["api"]
api = entry.runtime_data.api
try:
async with asyncio.timeout(10):
res = await api.request("GET")
json_data = await res.json()
hass.data[DOMAIN][entry.entry_id]["current"] = json_data
entry.runtime_data.current = json_data
return _preprocess(json_data)
except ClientResponseError as exc:
_LOGGER.warning("API fetch failed. Status: %s, - %s", exc.code, exc.message)
raise UpdateFailed(exc) from exc

hass.data[DOMAIN][entry.entry_id]["coordinator"] = DataUpdateCoordinator(
entry.runtime_data.coordinator = DataUpdateCoordinator(
hass,
logging.getLogger(__name__),
name=DOMAIN,
update_method=async_fetch,
update_interval=timedelta(minutes=5),
)
await hass.data[DOMAIN][entry.entry_id]["coordinator"].async_refresh()
return hass.data[DOMAIN][entry.entry_id]["coordinator"]
await entry.runtime_data.coordinator.async_refresh()
return entry.runtime_data.coordinator


async def async_migrate_entry(hass, config_entry: ConfigEntry):
Expand Down
64 changes: 32 additions & 32 deletions custom_components/weatherlink/binary_sensor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,18 @@
BinarySensorEntity,
BinarySensorEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity import DeviceInfo, EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from homeassistant.util import dt as dt_util

from . import SENSOR_TYPE_AIRLINK, SENSOR_TYPE_VUE_AND_VANTAGE_PRO, get_coordinator
from . import (
SENSOR_TYPE_AIRLINK,
SENSOR_TYPE_VUE_AND_VANTAGE_PRO,
WLConfigEntry,
get_coordinator,
)
from .const import (
CONF_API_VERSION,
CONFIG_URL,
Expand Down Expand Up @@ -70,17 +74,17 @@ class WLBinarySensorDescription(BinarySensorEntityDescription):

async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
entry: WLConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up the binary sensor platform."""
coordinator = await get_coordinator(hass, config_entry)
primary_tx_id = hass.data[DOMAIN][config_entry.entry_id]["primary_tx_id"]
coordinator = await get_coordinator(hass, entry)
primary_tx_id = entry.runtime_data.primary_tx_id

entities = [
WLSensor(coordinator, hass, config_entry, description, primary_tx_id)
WLSensor(coordinator, hass, entry, description, primary_tx_id)
for description in SENSOR_TYPES
if (config_entry.data[CONF_API_VERSION] not in description.exclude_api_ver)
if (entry.data[CONF_API_VERSION] not in description.exclude_api_ver)
and (
coordinator.data[primary_tx_id].get(DataKey.DATA_STRUCTURE)
not in description.exclude_data_structure
Expand All @@ -89,14 +93,14 @@ async def async_setup_entry(
]

aux_entities = []
if config_entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
for sensor in hass.data[DOMAIN][config_entry.entry_id]["sensors_metadata"]:
if entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
for sensor in entry.runtime_data.sensors_metadata:
if sensor["tx_id"] is not None and sensor["tx_id"] != primary_tx_id:
aux_entities += [
WLSensor(
coordinator,
hass,
config_entry,
entry,
description,
sensor["tx_id"],
)
Expand All @@ -110,7 +114,7 @@ async def async_setup_entry(
WLSensor(
coordinator,
hass,
config_entry,
entry,
description,
sensor["lsid"],
)
Expand All @@ -134,17 +138,17 @@ def __init__(
self,
coordinator,
hass: HomeAssistant,
entry: ConfigEntry,
entry: WLConfigEntry,
description: WLBinarySensorDescription,
tx_id: int,
):
"""Initialize the sensor."""
super().__init__(coordinator)
self.hass = hass
self.entry: ConfigEntry = entry
self.entry = entry
self.entity_description = description
self.tx_id = tx_id
self.primary_tx_id = self.hass.data[DOMAIN][entry.entry_id]["primary_tx_id"]
self.primary_tx_id = entry.runtime_data.primary_tx_id
self._attr_has_entity_name = True
tx_id_part = f"-{self.tx_id}" if self.tx_id != self.primary_tx_id else ""
self._attr_unique_id = (
Expand Down Expand Up @@ -175,17 +179,17 @@ def get_unique_id_base(self):
def get_firmware(self) -> str | None:
"""Get firmware version."""
if self.entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
return self.hass.data[DOMAIN][self.entry.entry_id]["station_data"][
"stations"
][0].get("firmware_version")
return self.entry.runtime_data.station_data["stations"][0].get(
"firmware_version"
)
return None

def get_serial(self) -> str | None:
"""Get serial number."""
if self.entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
return self.hass.data[DOMAIN][self.entry.entry_id]["station_data"][
"stations"
][0].get("gateway_id_hex")
return self.entry.runtime_data.station_data["stations"][0].get(
"gateway_id_hex"
)
return None

def generate_name(self):
Expand All @@ -194,12 +198,10 @@ def generate_name(self):
return self.coordinator.data[1]["station_name"]
if self.entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
if self.tx_id == self.primary_tx_id:
return self.hass.data[DOMAIN][self.entry.entry_id]["station_data"][
"stations"
][0]["station_name"]
for sensor in self.hass.data[DOMAIN][self.entry.entry_id][
"sensors_metadata"
]:
return self.entry.runtime_data.station_data["stations"][0][
"station_name"
]
for sensor in self.entry.runtime_data.sensors_metadata:
if sensor["sensor_type"] in (55, 56) and sensor["tx_id"] == self.tx_id:
return f"{sensor['product_name']} ID{sensor['tx_id']}"

Expand All @@ -216,14 +218,12 @@ def generate_model(self):
if self.entry.data[CONF_API_VERSION] == ApiVersion.API_V1:
return "WeatherLink - API V1"
if self.entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
model: str = self.hass.data[DOMAIN][self.entry.entry_id]["station_data"][
"stations"
][0].get("product_number")
model: str = self.entry.runtime_data.station_data["stations"][0].get(
"product_number"
)
break_out = False
product_name = ""
for sensor in self.hass.data[DOMAIN][self.entry.entry_id][
"sensors_metadata"
]:
for sensor in self.entry.runtime_data.sensors_metadata:
if break_out:
break
if (
Expand Down
24 changes: 10 additions & 14 deletions custom_components/weatherlink/diagnostics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
from __future__ import annotations

from homeassistant.components.diagnostics import async_redact_data
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
Expand All @@ -13,10 +12,11 @@
CONF_API_SECRET,
CONF_API_TOKEN,
CONF_API_VERSION,
DOMAIN,
ApiVersion,
)

from . import WLConfigEntry

TO_REDACT = {
CONF_PASSWORD,
CONF_USERNAME,
Expand All @@ -28,24 +28,20 @@


async def async_get_config_entry_diagnostics(
hass: HomeAssistant, config_entry: ConfigEntry
hass: HomeAssistant, entry: WLConfigEntry
) -> dict:
"""Return diagnostics for a config entry."""
coordinator: DataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id][
"coordinator"
]
station_data = hass.data[DOMAIN][config_entry.entry_id].get("station_data", {})
current = hass.data[DOMAIN][config_entry.entry_id]["current"]
sensor_metadata = hass.data[DOMAIN][config_entry.entry_id]["sensors_metadata"]
coordinator: DataUpdateCoordinator = entry.runtime_data.coordinator
station_data = entry.runtime_data.station_data
current = entry.runtime_data.current
sensor_metadata = entry.runtime_data.sensors_metadata

sensor_data = {}
if config_entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
sensor_data = await hass.data[DOMAIN][config_entry.entry_id][
"api"
].get_all_sensors()
if entry.data[CONF_API_VERSION] == ApiVersion.API_V2:
sensor_data = await entry.runtime_data.api.get_all_sensors()

diagnostics_data = {
"info": async_redact_data(config_entry.data, TO_REDACT),
"info": async_redact_data(entry.data, TO_REDACT),
"station_data": async_redact_data(station_data, TO_REDACT),
"all_sensor_data": async_redact_data(sensor_data, TO_REDACT),
"sensor_metadata": async_redact_data(sensor_metadata, TO_REDACT),
Expand Down
Loading

0 comments on commit 5977bc6

Please sign in to comment.