diff --git a/custom_components/load_shedding/__init__.py b/custom_components/load_shedding/__init__.py index ccce85d..23eb2f7 100644 --- a/custom_components/load_shedding/__init__.py +++ b/custom_components/load_shedding/__init__.py @@ -1,12 +1,13 @@ """The LoadShedding component.""" from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone import logging from typing import Any from load_shedding.libs.sepush import SePush, SePushError from load_shedding.providers import Area, Stage +import urllib3 from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( @@ -186,7 +187,7 @@ def __init__(self, hass: HomeAssistant, sepush: SePush) -> None: async def _async_update_data(self) -> dict: """Retrieve latest load shedding data.""" - now = datetime.now(datetime.UTC).replace(microsecond=0) + now = datetime.now(UTC).replace(microsecond=0) diff = 0 if self.last_update is not None: diff = (now - self.last_update).seconds @@ -196,6 +197,9 @@ async def _async_update_data(self) -> dict: try: stage = await self.async_update_stage() + except SePushError as err: + _LOGGER.error("Unable to get stage: %s", err) + self.data = {} except UpdateFailed as err: _LOGGER.exception("Unable to get stage: %s", err) self.data = {} @@ -207,60 +211,55 @@ async def _async_update_data(self) -> dict: async def async_update_stage(self) -> dict: """Retrieve latest stage.""" - now = datetime.now(datetime.UTC).replace(microsecond=0) - try: - esp = await self.hass.async_add_executor_job(self.sepush.status) - except SePushError as err: - raise UpdateFailed(err) from err - else: - data = {} - statuses = esp.get("status", {}) - for idx, area in statuses.items(): - stage = Stage(int(area.get("stage", "0"))) - start_time = datetime.fromisoformat(area.get("stage_updated")) + now = datetime.now(UTC).replace(microsecond=0) + esp = await self.hass.async_add_executor_job(self.sepush.status) + + data = {} + statuses = esp.get("status", {}) + for idx, area in statuses.items(): + stage = Stage(int(area.get("stage", "0"))) + start_time = datetime.fromisoformat(area.get("stage_updated")) + start_time = start_time.replace(second=0, microsecond=0) + planned = [ + { + ATTR_STAGE: stage, + ATTR_START_TIME: start_time.astimezone(UTC), + } + ] + + next_stages = area.get("next_stages", []) + for i, next_stage in enumerate(next_stages): + # Prev + prev_end = datetime.fromisoformat( + next_stage.get("stage_start_timestamp") + ) + prev_end = prev_end.replace(second=0, microsecond=0) + planned[i][ATTR_END_TIME] = prev_end.astimezone(UTC) + + # Next + stage = Stage(int(next_stage.get("stage", "0"))) + start_time = datetime.fromisoformat( + next_stage.get("stage_start_timestamp") + ) start_time = start_time.replace(second=0, microsecond=0) - planned = [ + planned.append( { ATTR_STAGE: stage, - ATTR_START_TIME: start_time.astimezone(datetime.UTC), + ATTR_START_TIME: start_time.astimezone(UTC), } - ] - - next_stages = area.get("next_stages", []) - for i, next_stage in enumerate(next_stages): - # Prev - prev_end = datetime.fromisoformat( - next_stage.get("stage_start_timestamp") - ) - prev_end = prev_end.replace(second=0, microsecond=0) - planned[i][ATTR_END_TIME] = prev_end.astimezone(datetime.UTC) - - # Next - stage = Stage(int(next_stage.get("stage", "0"))) - start_time = datetime.fromisoformat( - next_stage.get("stage_start_timestamp") - ) - start_time = start_time.replace(second=0, microsecond=0) - planned.append( - { - ATTR_STAGE: stage, - ATTR_START_TIME: start_time.astimezone(datetime.UTC), - } - ) + ) - filtered = [] - for stage in planned: - if ATTR_END_TIME not in stage: - stage[ATTR_END_TIME] = stage[ATTR_START_TIME] + timedelta( - days=7 - ) - if ATTR_END_TIME in stage and stage.get(ATTR_END_TIME) >= now: - filtered.append(stage) + filtered = [] + for stage in planned: + if ATTR_END_TIME not in stage: + stage[ATTR_END_TIME] = stage[ATTR_START_TIME] + timedelta(days=7) + if ATTR_END_TIME in stage and stage.get(ATTR_END_TIME) >= now: + filtered.append(stage) - data[idx] = { - ATTR_NAME: area.get("name", ""), - ATTR_PLANNED: filtered, - } + data[idx] = { + ATTR_NAME: area.get("name", ""), + ATTR_PLANNED: filtered, + } return data @@ -289,7 +288,7 @@ def add_area(self, area: Area = None) -> None: async def _async_update_data(self) -> dict: """Retrieve latest load shedding data.""" - now = datetime.now(datetime.UTC).replace(microsecond=0) + now = datetime.now(UTC).replace(microsecond=0) diff = 0 if self.last_update is not None: diff = (now - self.last_update).seconds @@ -300,6 +299,9 @@ async def _async_update_data(self) -> dict: try: area = await self.async_update_area() + except SePushError as err: + _LOGGER.error("Unable to get area schedule: %s", err) + self.data = {} except UpdateFailed as err: _LOGGER.exception("Unable to get area schedule: %s", err) self.data = {} @@ -315,10 +317,7 @@ async def async_update_area(self) -> dict: area_id_data: dict = {} for area in self.areas: - try: - esp = await self.hass.async_add_executor_job(self.sepush.area, area.id) - except SePushError as err: - raise UpdateFailed(err) from err + esp = await self.hass.async_add_executor_job(self.sepush.area, area.id) # Get events for area events = [] @@ -332,10 +331,8 @@ async def async_update_area(self) -> dict: if note == str(Stage.LOAD_REDUCTION): stage = Stage.LOAD_REDUCTION - start = datetime.fromisoformat(event.get("start")).astimezone( - datetime.UTC - ) - end = datetime.fromisoformat(event.get("end")).astimezone(datetime.UTC) + start = datetime.fromisoformat(event.get("start")).astimezone(UTC) + end = datetime.fromisoformat(event.get("end")).astimezone(UTC) events.append( { @@ -478,7 +475,7 @@ def utc_dt(date: datetime, time: datetime) -> datetime: second=0, microsecond=0, tzinfo=sast, - ).astimezone(datetime.UTC) + ).astimezone(UTC) class LoadSheddingQuotaCoordinator(DataUpdateCoordinator[dict[str, Any]]): @@ -494,9 +491,12 @@ def __init__(self, hass: HomeAssistant, sepush: SePush) -> None: async def _async_update_data(self) -> dict: """Retrieve latest load shedding data.""" - now = datetime.now(datetime.UTC).replace(microsecond=0) + now = datetime.now(UTC).replace(microsecond=0) try: quota = await self.async_update_quota() + except SePushError as err: + _LOGGER.error("Unable to get quota: %s", err) + self.data = {} except UpdateFailed as err: _LOGGER.exception("Unable to get quota: %s", err) else: @@ -507,10 +507,7 @@ async def _async_update_data(self) -> dict: async def async_update_quota(self) -> dict: """Retrieve latest quota.""" - try: - esp = await self.hass.async_add_executor_job(self.sepush.check_allowance) - except SePushError as err: - raise UpdateFailed(err) from err + esp = await self.hass.async_add_executor_job(self.sepush.check_allowance) return esp.get("allowance", {}) diff --git a/custom_components/load_shedding/sensor.py b/custom_components/load_shedding/sensor.py index 9b5fce4..d99ef40 100644 --- a/custom_components/load_shedding/sensor.py +++ b/custom_components/load_shedding/sensor.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass -from datetime import datetime, timedelta +from datetime import UTC, datetime, timedelta from typing import Any, cast from load_shedding.providers import Area, Stage @@ -153,7 +153,7 @@ def extra_state_attributes(self) -> dict[str, list, Any]: if not self.data: return self._attr_extra_state_attributes - now = datetime.now(datetime.UTC) + now = datetime.now(UTC) # data = get_sensor_attrs(planned, planned[0].get(ATTR_STAGE, Stage.UNKNOWN)) # data[ATTR_PLANNED] = [] data = dict(self._attr_extra_state_attributes) @@ -200,8 +200,6 @@ class LoadSheddingAreaSensorEntity( ): """Define a LoadShedding Area sensor entity.""" - coordinator: CoordinatorEntity - def __init__(self, coordinator: CoordinatorEntity, area: Area) -> None: """Initialize.""" super().__init__(coordinator) @@ -241,7 +239,7 @@ def native_value(self) -> StateType: if not events: return STATE_OFF - now = datetime.now(datetime.UTC) + now = datetime.now(UTC) for event in events: if ATTR_END_TIME in event and event.get(ATTR_END_TIME) < now: @@ -270,7 +268,7 @@ def extra_state_attributes(self) -> dict[str, list, Any]: if not self.data: return self._attr_extra_state_attributes - now = datetime.now(datetime.UTC) + now = datetime.now(UTC) data = dict(self._attr_extra_state_attributes) if events := self.data.get(ATTR_FORECAST, []): data[ATTR_FORECAST] = [] @@ -388,7 +386,7 @@ def get_sensor_attrs(forecast: list, stage: Stage = Stage.NO_LOAD_SHEDDING) -> d ATTR_STAGE: stage.value, } - now = datetime.now(datetime.UTC) + now = datetime.now(UTC) data = dict(DEFAULT_DATA) data[ATTR_STAGE] = stage.value