Refactor and use latest observation as radar thumbnail

jdejaegh · jdejaegh · commit ed20cd99221e · 2023-12-28T15:30:30.000+01:00
diff --git a/custom_components/irm_kmi/camera.py b/custom_components/irm_kmi/camera.py
@@ -2,7 +2,6 @@
 # File inspired by https://github.com/jodur/imagesdirectory-camera/blob/main/custom_components/imagedirectory/camera.py
 
 import logging
-import os
 
 from homeassistant.components.camera import Camera, async_get_still_stream
 from homeassistant.config_entries import ConfigEntry
@@ -22,7 +21,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e
 
     _LOGGER.debug(f'async_setup_entry entry is: {entry}')
     coordinator = hass.data[DOMAIN][entry.entry_id]
-    # await coordinator.async_config_entry_first_refresh()
     async_add_entities(
         [IrmKmiRadar(coordinator, entry)]
     )
@@ -57,10 +55,7 @@ def frame_interval(self) -> float:
     def camera_image(self,
                      width: int | None = None,
                      height: int | None = None) -> bytes | None:
-        images = self.coordinator.data.get('animation', {}).get('images')
-        if isinstance(images, list) and len(images) > 0:
-            return images[0]
-        return None
+        return self.coordinator.data.get('animation', {}).get('most_recent_image')
 
     async def async_camera_image(
             self,
@@ -84,10 +79,10 @@ async def handle_async_mjpeg_stream(self, request):
         return await self.handle_async_still_stream(request, self.frame_interval)
 
     async def iterate(self) -> bytes | None:
-        images = self.coordinator.data.get('animation', {}).get('images')
-        if isinstance(images, list) and len(images) > 0:
-            r = images[self._image_index]
-            self._image_index = (self._image_index + 1) % len(images)
+        sequence = self.coordinator.data.get('animation', {}).get('sequence')
+        if isinstance(sequence, list) and len(sequence) > 0:
+            r = sequence[self._image_index].get('image', None)
+            self._image_index = (self._image_index + 1) % len(sequence)
             return r
         return None
 
diff --git a/custom_components/irm_kmi/coordinator.py b/custom_components/irm_kmi/coordinator.py
@@ -6,6 +6,7 @@
 from typing import List
 
 import async_timeout
+import pytz
 from homeassistant.components.weather import Forecast
 from homeassistant.const import ATTR_LATITUDE, ATTR_LONGITUDE
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -16,7 +17,8 @@
 from .api import IrmKmiApiClient, IrmKmiApiError
 from .const import IRM_KMI_TO_HA_CONDITION_MAP as CDT_MAP
 from .const import OUT_OF_BENELUX
-from .data import IrmKmiForecast
+from .data import (AnimationFrameData, CurrentWeatherData, IrmKmiForecast,
+                   ProcessedCoordinatorData, RadarAnimationData)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -37,7 +39,7 @@ def __init__(self, hass, zone):
         self._api_client = IrmKmiApiClient(session=async_get_clientsession(hass))
         self._zone = zone
 
-    async def _async_update_data(self):
+    async def _async_update_data(self) -> ProcessedCoordinatorData:
         """Fetch data from API endpoint.
 
         This is the place to pre-process the data to lookup tables
@@ -62,67 +64,107 @@ async def _async_update_data(self):
         if api_data.get('cityName', None) in OUT_OF_BENELUX:
             raise UpdateFailed(f"Zone '{self._zone}' is out of Benelux and forecast is only available in the Benelux")
 
-        result = self.process_api_data(api_data)
+        return await self.process_api_data(api_data)
 
-        # TODO make such that the most up to date image is specified to entity for static display
-        return result | await self._async_animation_data(api_data)
+    async def _async_animation_data(self, api_data: dict) -> RadarAnimationData:
 
-    async def _async_animation_data(self, api_data: dict) -> dict:
-
-        default = {'animation': None}
         animation_data = api_data.get('animation', {}).get('sequence')
-        localisation_layer = api_data.get('animation', {}).get('localisationLayer')
-        country = api_data.get('country', None)
+        localisation_layer_url = api_data.get('animation', {}).get('localisationLayer')
+        country = api_data.get('country', '')
+
+        if animation_data is None or localisation_layer_url is None or not isinstance(animation_data, list):
+            return RadarAnimationData()
+
+        try:
+            images_from_api = await self.download_images_from_api(animation_data, country, localisation_layer_url)
+        except IrmKmiApiError:
+            _LOGGER.warning(f"Could not get images for weather radar")
+            return RadarAnimationData()
+
+        localisation = Image.open(BytesIO(images_from_api[0])).convert('RGBA')
+        images_from_api = images_from_api[1:]
 
-        if animation_data is None or localisation_layer is None or not isinstance(animation_data, list):
-            return default
+        radar_animation = await self.merge_frames_from_api(animation_data, country, images_from_api, localisation)
+        # TODO support translation here
+        radar_animation['hint'] = api_data.get('animation', {}).get('sequenceHint', {}).get('en')
+        return radar_animation
 
+    async def download_images_from_api(self, animation_data, country, localisation_layer_url):
         coroutines = list()
-        coroutines.append(self._api_client.get_image(f"{localisation_layer}&th={'d' if country == 'NL' else 'n'}"))
+        coroutines.append(self._api_client.get_image(f"{localisation_layer_url}&th={'d' if country == 'NL' else 'n'}"))
         for frame in animation_data:
             if frame.get('uri', None) is not None:
                 coroutines.append(self._api_client.get_image(frame.get('uri')))
+        async with async_timeout.timeout(20):
+            images_from_api = await asyncio.gather(*coroutines, return_exceptions=True)
 
-        try:
-            async with async_timeout.timeout(20):
-                r = await asyncio.gather(*coroutines, return_exceptions=True)
-        except IrmKmiApiError:
-            _LOGGER.warning(f"Could not get images for weather radar")
-            return default
-        _LOGGER.debug(f"Just downloaded {len(r)} images")
+        _LOGGER.debug(f"Just downloaded {len(images_from_api)} images")
+        return images_from_api
+
+    async def merge_frames_from_api(self, animation_data, country, images_from_api,
+                                    localisation_layer) -> RadarAnimationData:
 
         if country == 'NL':
             background = Image.open("custom_components/irm_kmi/resources/nl.png").convert('RGBA')
         else:
             background = Image.open("custom_components/irm_kmi/resources/be_bw.png").convert('RGBA')
-        localisation = Image.open(BytesIO(r[0])).convert('RGBA')
-        merged_frames = list()
-        for frame in r[1:]:
+
+        most_recent_frame = None
+        tz = pytz.timezone(self.hass.config.time_zone)
+        current_time = datetime.now(tz=tz)
+        sequence: List[AnimationFrameData] = list()
+        for (idx, sequence_element) in enumerate(animation_data):
+            frame = images_from_api[idx]
             layer = Image.open(BytesIO(frame)).convert('RGBA')
             temp = Image.alpha_composite(background, layer)
-            temp = Image.alpha_composite(temp, localisation)
+            temp = Image.alpha_composite(temp, localisation_layer)
 
             draw = ImageDraw.Draw(temp)
             font = ImageFont.truetype("custom_components/irm_kmi/resources/roboto_medium.ttf", 16)
-            # TODO write actual date time
+            time_image = (datetime.fromisoformat(sequence_element.get('time'))
+                          .astimezone(tz=tz))
+
+            time_str = time_image.isoformat(sep=' ', timespec='minutes')
+
             if country == 'NL':
-                draw.text((4, 4), "Sample Text", (0, 0, 0), font=font)
+                draw.text((4, 4), time_str, (0, 0, 0), font=font)
             else:
-                draw.text((4, 4), "Sample Text", (255, 255, 255), font=font)
+                draw.text((4, 4), time_str, (255, 255, 255), font=font)
 
             bytes_img = BytesIO()
-            temp.save(bytes_img, 'png')
-            merged_frames.append(bytes_img.getvalue())
+            temp.save(bytes_img, 'png', compress_level=8)
+
+            sequence.append(
+                AnimationFrameData(
+                    time=time_image,
+                    image=bytes_img.getvalue()
+                )
+            )
+
+            if most_recent_frame is None and current_time < time_image:
+                recent_idx = idx - 1 if idx > 0 else idx
+                most_recent_frame = sequence[recent_idx].get('image', None)
+                _LOGGER.debug(f"Most recent frame is at {sequence[recent_idx].get('time')}")
+
+        background.close()
+        most_recent_frame = most_recent_frame if most_recent_frame is not None else sequence[-1].get('image')
 
-        return {'animation': {
-            'images': merged_frames,
-            # TODO support translation for hint
-            'hint': api_data.get('animation', {}).get('sequenceHint', {}).get('en')
-        }
-        }
+        return RadarAnimationData(
+            sequence=sequence,
+            most_recent_image=most_recent_frame
+        )
+
+    async def process_api_data(self, api_data: dict) -> ProcessedCoordinatorData:
+
+        return ProcessedCoordinatorData(
+            current_weather=IrmKmiCoordinator.current_weather_from_data(api_data),
+            daily_forecast=IrmKmiCoordinator.daily_list_to_forecast(api_data.get('for', {}).get('daily')),
+            hourly_forecast=IrmKmiCoordinator.hourly_list_to_forecast(api_data.get('for', {}).get('hourly')),
+            animation=await self._async_animation_data(api_data=api_data)
+        )
 
     @staticmethod
-    def process_api_data(api_data):
+    def current_weather_from_data(api_data: dict) -> CurrentWeatherData:
         # Process data to get current hour forecast
         now_hourly = None
         hourly_forecast_data = api_data.get('for', {}).get('hourly')
@@ -140,23 +182,18 @@ def process_api_data(api_data):
             for module in module_data:
                 if module.get('type', None) == 'uv':
                     uv_index = module.get('data', {}).get('levelValue')
-        # Put everything together
+
         # TODO NL cities have a better 'obs' section, use that for current weather
-        processed_data = {
-            'current_weather': {
-                'condition': CDT_MAP.get(
-                    (api_data.get('obs', {}).get('ww'), api_data.get('obs', {}).get('dayNight')), None),
-                'temperature': api_data.get('obs', {}).get('temp'),
-                'wind_speed': now_hourly.get('windSpeedKm', None) if now_hourly is not None else None,
-                'wind_gust_speed': now_hourly.get('windPeakSpeedKm', None) if now_hourly is not None else None,
-                'wind_bearing': now_hourly.get('windDirectionText', {}).get('en') if now_hourly is not None else None,
-                'pressure': now_hourly.get('pressure', None) if now_hourly is not None else None,
-                'uv_index': uv_index
-            },
-            'daily_forecast': IrmKmiCoordinator.daily_list_to_forecast(api_data.get('for', {}).get('daily')),
-            'hourly_forecast': IrmKmiCoordinator.hourly_list_to_forecast(api_data.get('for', {}).get('hourly'))
-        }
-        return processed_data
+        current_weather = CurrentWeatherData(
+            condition=CDT_MAP.get((api_data.get('obs', {}).get('ww'), api_data.get('obs', {}).get('dayNight')), None),
+            temperature=api_data.get('obs', {}).get('temp'),
+            wind_speed=now_hourly.get('windSpeedKm', None) if now_hourly is not None else None,
+            wind_gust_speed=now_hourly.get('windPeakSpeedKm', None) if now_hourly is not None else None,
+            wind_bearing=now_hourly.get('windDirectionText', {}).get('en') if now_hourly is not None else None,
+            pressure=now_hourly.get('pressure', None) if now_hourly is not None else None,
+            uv_index=uv_index
+        )
+        return current_weather
 
     @staticmethod
     def hourly_list_to_forecast(data: List[dict] | None) -> List[Forecast] | None:
diff --git a/custom_components/irm_kmi/data.py b/custom_components/irm_kmi/data.py
@@ -1,4 +1,7 @@
 """Data classes for IRM KMI integration"""
+from datetime import datetime
+from typing import List, TypedDict
+
 from homeassistant.components.weather import Forecast
 
 
@@ -8,3 +11,31 @@ class IrmKmiForecast(Forecast):
     # TODO: add condition_2 as well and evolution to match data from the API?
     text_fr: str | None
     text_nl: str | None
+
+
+class CurrentWeatherData(TypedDict, total=False):
+    condition: str | None
+    temperature: float | None
+    wind_speed: float | None
+    wind_gust_speed: float | None
+    wind_bearing: float | str | None
+    uv_index: float | None
+    pressure: float | None
+
+
+class AnimationFrameData(TypedDict, total=False):
+    time: datetime | None
+    image: bytes | None
+
+
+class RadarAnimationData(TypedDict, total=False):
+    sequence: List[AnimationFrameData] | None
+    most_recent_image: bytes | None
+    hint: str | None
+
+
+class ProcessedCoordinatorData(TypedDict, total=False):
+    current_weather: CurrentWeatherData
+    hourly_forecast: List[Forecast] | None
+    daily_forecast: List[IrmKmiForecast] | None
+    animation: RadarAnimationData
diff --git a/custom_components/irm_kmi/weather.py b/custom_components/irm_kmi/weather.py
@@ -24,7 +24,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e
 
     _LOGGER.debug(f'async_setup_entry entry is: {entry}')
     coordinator = hass.data[DOMAIN][entry.entry_id]
-    # await coordinator.async_config_entry_first_refresh()
     async_add_entities(
         [IrmKmiWeather(coordinator, entry)]
     )
diff --git a/requirements.txt b/requirements.txt
@@ -2,4 +2,5 @@ aiohttp==3.9.1
 async_timeout==4.0.3
 homeassistant==2023.12.3
 voluptuous==0.13.1
-Pillow==10.1.0
+Pillow==10.1.0
+pytz==2023.3.post1
diff --git a/tests/test_coordinator.py b/tests/test_coordinator.py
@@ -8,7 +8,7 @@
 from pytest_homeassistant_custom_component.common import load_fixture
 
 from custom_components.irm_kmi.coordinator import IrmKmiCoordinator
-from custom_components.irm_kmi.data import IrmKmiForecast
+from custom_components.irm_kmi.data import CurrentWeatherData, IrmKmiForecast
 
 
 def get_api_data() -> dict:
@@ -19,17 +19,17 @@ def get_api_data() -> dict:
 @freeze_time(datetime.fromisoformat('2023-12-26T18:30:00.028724'))
 def test_current_weather() -> None:
     api_data = get_api_data()
-    result = IrmKmiCoordinator.process_api_data(api_data).get('current_weather')
-
-    expected = {
-        'condition': ATTR_CONDITION_CLOUDY,
-        'temperature': 7,
-        'wind_speed': 5,
-        'wind_gust_speed': None,
-        'wind_bearing': 'WSW',
-        'pressure': 1020,
-        'uv_index': .7
-    }
+    result = IrmKmiCoordinator.current_weather_from_data(api_data)
+
+    expected = CurrentWeatherData(
+        condition=ATTR_CONDITION_CLOUDY,
+        temperature=7,
+        wind_speed=5,
+        wind_gust_speed=None,
+        wind_bearing='WSW',
+        pressure=1020,
+        uv_index=.7
+    )
 
     assert result == expected
 
@@ -83,4 +83,3 @@ def test_hourly_forecast() -> None:
     )
 
     assert result[8] == expected
-

Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e`
`24`	`24`
`25`	`25`	`_LOGGER.debug(f'async_setup_entry entry is: {entry}')`
`26`	`26`	`coordinator = hass.data[DOMAIN][entry.entry_id]`
`27`		`- # await coordinator.async_config_entry_first_refresh()`
`28`	`27`	`async_add_entities(`
`29`	`28`	`[IrmKmiWeather(coordinator, entry)]`
`30`	`29`	`)`