From f531a383315296ec95042f4ad6a28dc7f20a7c19 Mon Sep 17 00:00:00 2001 From: balgerion <133121849+balgerion@users.noreply.github.com> Date: Mon, 28 Apr 2025 11:19:55 +0200 Subject: [PATCH] Update sensor.py --- custom_components/pstryk/sensor.py | 310 ++++++++--------------------- 1 file changed, 87 insertions(+), 223 deletions(-) diff --git a/custom_components/pstryk/sensor.py b/custom_components/pstryk/sensor.py index 68fc509..875cf2a 100644 --- a/custom_components/pstryk/sensor.py +++ b/custom_components/pstryk/sensor.py @@ -1,236 +1,100 @@ import logging -from datetime import timedelta -from homeassistant.components.sensor import SensorEntity -from homeassistant.helpers.event import async_track_point_in_time +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.components.sensor import SensorEntity, SensorStateClass +from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import dt as dt_util -import aiohttp -import async_timeout -import asyncio - -from .const import API_URL, DOMAIN +from .update_coordinator import PstrykDataUpdateCoordinator +from .const import DOMAIN _LOGGER = logging.getLogger(__name__) -SCAN_INTERVAL = timedelta(minutes=30) +async def async_setup_entry( + hass: HomeAssistant, + entry: ConfigEntry, + async_add_entities, +) -> None: + """Set up the four Pstryk sensors via the coordinator.""" + api_key = hass.data[DOMAIN][entry.entry_id]["api_key"] + buy_top = entry.options.get("buy_top", entry.data.get("buy_top", 5)) + sell_top = entry.options.get("sell_top", entry.data.get("sell_top", 5)) -def convert_price(value): - try: - return float(str(value).replace(",", ".").strip()) - except (ValueError, TypeError) as e: - _LOGGER.warning("Price conversion error: %s", str(e)) - return None + entities = [] + for price_type in ("buy", "sell"): + key = f"{entry.entry_id}_{price_type}" + coordinator: PstrykDataUpdateCoordinator = hass.data[DOMAIN].get(key) + if not coordinator: + coordinator = PstrykDataUpdateCoordinator(hass, api_key, price_type) + await coordinator.async_config_entry_first_refresh() + coordinator.schedule_hourly_update() + coordinator.schedule_midnight_update() + hass.data[DOMAIN][key] = coordinator -async def async_setup_entry(hass, config_entry, async_add_entities): - config = config_entry.data - sensors = [ - PstrykCurrentPriceSensor(config["api_key"], "buy"), - PstrykCurrentPriceSensor(config["api_key"], "sell"), - PstrykPriceTableSensor(config["api_key"], "buy", config.get("buy_top", 5)), - PstrykPriceTableSensor(config["api_key"], "sell", config.get("sell_top", 5)) - ] - async_add_entities(sensors, True) + entities.append(PstrykCurrentPriceSensor(coordinator, price_type)) + top = buy_top if price_type == "buy" else sell_top + entities.append(PstrykPriceTableSensor(coordinator, price_type, top)) -class PstrykCurrentPriceSensor(SensorEntity): - def __init__(self, api_key, price_type): - self._api_key = api_key - self._price_type = price_type - self._state = None - self._available = True - self._unsub_update = None + async_add_entities(entities, True) + + +class PstrykCurrentPriceSensor(CoordinatorEntity, SensorEntity): + """Current price sensor.""" + _attr_state_class = SensorStateClass.MEASUREMENT + + def __init__(self, coordinator: PstrykDataUpdateCoordinator, price_type: str): + super().__init__(coordinator) + self.price_type = price_type @property - def name(self): return f"Pstryk Current {self._price_type.title()} Price" - - @property - def unique_id(self): return f"pstryk_current_{self._price_type}_price" - - @property - def native_value(self): return self._state - - @property - def native_unit_of_measurement(self): return "PLN/kWh" - - @property - def available(self): return self._available + def name(self) -> str: + return f"Pstryk Current {self.price_type.title()} Price" - async def async_added_to_hass(self): - await super().async_added_to_hass() - self._schedule_next_update() + @property + def unique_id(self) -> str: + return f"{DOMAIN}_{self.price_type}_current" - def _schedule_next_update(self): - if self._unsub_update: - self._unsub_update() - - now_local = dt_util.now() - next_midnight_local = (now_local + timedelta(days=1)).replace( - hour=0, minute=1, second=0, microsecond=0 + @property + def native_value(self): + return self.coordinator.data.get("current") + + @property + def native_unit_of_measurement(self) -> str: + return "PLN/kWh" + + +class PstrykPriceTableSensor(CoordinatorEntity, SensorEntity): + """Today's price table sensor.""" + _attr_state_class = SensorStateClass.MEASUREMENT + + def __init__(self, coordinator: PstrykDataUpdateCoordinator, price_type: str, top_count: int): + super().__init__(coordinator) + self.price_type = price_type + self.top_count = top_count + + @property + def name(self) -> str: + return f"Pstryk {self.price_type.title()} Price Table" + + @property + def unique_id(self) -> str: + return f"{DOMAIN}_{self.price_type}_table" + + @property + def native_value(self) -> int: + # number of price slots today + return len(self.coordinator.data.get("prices_today", [])) + + @property + def extra_state_attributes(self) -> dict: + today = self.coordinator.data.get("prices_today", []) + sorted_prices = sorted( + today, + key=lambda x: x["price"], + reverse=(self.price_type == "sell"), ) - next_midnight_utc = dt_util.as_utc(next_midnight_local) - - self._unsub_update = async_track_point_in_time( - self.hass, - self._async_update_task, - next_midnight_utc - ) - - _LOGGER.debug( - "Następna aktualizacja %s: %s (Twoja strefa)", - self._price_type, - dt_util.as_local(next_midnight_utc).strftime("%Y-%m-%d %H:%M:%S") - ) - - async def _async_update_task(self, _): - try: - await self.async_update(no_throttle=True) - self.async_write_ha_state() - finally: - self._schedule_next_update() - - async def async_update(self, **kwargs): - try: - today_local = dt_util.now().replace(hour=0, minute=0, second=0, microsecond=0) - window_end_local = today_local + timedelta(days=1) - today_utc = dt_util.as_utc(today_local) - window_end_utc = dt_util.as_utc(window_end_local) - - endpoint = "pricing/" if self._price_type == "buy" else "prosumer-pricing/" - url = f"{API_URL}{endpoint}?resolution=hour&window_start={today_utc.strftime('%Y-%m-%dT%H:%M:%SZ')}&window_end={window_end_utc.strftime('%Y-%m-%dT%H:%M:%SZ')}" - - async with aiohttp.ClientSession() as session: - response = await session.get(url, headers={"Authorization": self._api_key, "Accept": "application/json"}) - data = await response.json() - - now = dt_util.utcnow() - current_price = None - for frame in data.get("frames", []): - start = dt_util.parse_datetime(frame["start"]) - end = dt_util.parse_datetime(frame["end"]) - if start <= now < end: - current_price = convert_price(frame["price_gross"]) - break - - self._state = current_price - self._available = current_price is not None - _LOGGER.debug( - "Zaktualizowano %s: %s PLN (Twoja strefa: %s)", - self._price_type, - self._state, - dt_util.as_local(now).strftime("%Y-%m-%d %H:%M:%S") - ) - - except Exception as e: - _LOGGER.error("Błąd aktualizacji %s: %s", self._price_type, str(e)) - self._state = None - self._available = False - - async def async_will_remove_from_hass(self): - if self._unsub_update: - self._unsub_update() - -class PstrykPriceTableSensor(SensorEntity): - def __init__(self, api_key, price_type, top_count): - self._api_key = api_key - self._price_type = price_type - self._top_count = top_count - self._state = None - self._attributes = {} - self._available = True - self._unsub_update = None - - @property - def name(self): return f"Pstryk {self._price_type.title()} Price Table" - - @property - def unique_id(self): return f"pstryk_{self._price_type}_price_table" - - @property - def native_value(self): return self._state - - @property - def extra_state_attributes(self): return self._attributes - - @property - def available(self): return self._available - - async def async_added_to_hass(self): - await super().async_added_to_hass() - self._schedule_next_update() - - def _schedule_next_update(self): - if self._unsub_update: - self._unsub_update() - - now_local = dt_util.now() - next_midnight_local = (now_local + timedelta(days=1)).replace( - hour=0, minute=1, second=0, microsecond=0 - ) - next_midnight_utc = dt_util.as_utc(next_midnight_local) - - self._unsub_update = async_track_point_in_time( - self.hass, - self._async_update_task, - next_midnight_utc - ) - - async def _async_update_task(self, _): - try: - await self.async_update(no_throttle=True) - self.async_write_ha_state() - finally: - self._schedule_next_update() - - def _convert_time(self, utc_str): - try: - dt_utc = dt_util.parse_datetime(utc_str) - dt_local = dt_util.as_local(dt_utc) - return dt_local.strftime("%Y-%m-%d %H:%M:%S") - except Exception as e: - _LOGGER.error("Błąd konwersji czasu: %s", str(e)) - return "N/A" - - async def async_update(self, **kwargs): - try: - today_local = dt_util.now().replace(hour=0, minute=0, second=0, microsecond=0) - window_end_local = today_local + timedelta(days=1) - today_utc = dt_util.as_utc(today_local) - window_end_utc = dt_util.as_utc(window_end_local) - - endpoint = "pricing/" if self._price_type == "buy" else "prosumer-pricing/" - url = f"{API_URL}{endpoint}?resolution=hour&window_start={today_utc.strftime('%Y-%m-%dT%H:%M:%SZ')}&window_end={window_end_utc.strftime('%Y-%m-%dT%H:%M:%SZ')}" - - async with aiohttp.ClientSession() as session: - response = await session.get(url, headers={"Authorization": self._api_key, "Accept": "application/json"}) - data = await response.json() - - prices = [] - for frame in data.get("frames", []): - price = convert_price(frame.get("price_gross")) - if price is not None: - prices.append({ - "start": self._convert_time(frame["start"]), - "price": price - }) - - sorted_prices = sorted(prices, key=lambda x: x["price"], reverse=(self._price_type == "sell")) - self._state = len(prices) - self._attributes = { - "all_prices": [{"start": p["start"], "price": p["price"]} for p in prices], - "best_prices": [{"start": p["start"], "price": p["price"]} for p in sorted_prices[:self._top_count]], - "top_count": self._top_count, - "last_updated": dt_util.as_local(dt_util.utcnow()).isoformat() - } - self._available = True - _LOGGER.debug( - "Zaktualizowano tabelę %s (Twoja strefa: %s)", - self._price_type, - dt_util.as_local(dt_util.utcnow()).strftime("%Y-%m-%d %H:%M:%S") - ) - - except Exception as e: - _LOGGER.error("Błąd aktualizacji tabeli %s: %s", self._price_type, str(e)) - self._state = None - self._available = False - - async def async_will_remove_from_hass(self): - if self._unsub_update: - self._unsub_update() + return { + "all_prices": today, + "best_prices": sorted_prices[: self.top_count], + "top_count": self.top_count, + "last_updated": dt_util.as_local(dt_util.utcnow()).isoformat(), + }