bump version 1.0.3 -> 1.1.0

Add support for heat warning

.github/workflows/pytest.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.11", "3.13"]
+        python-version: ["3.12", "3.13"]
 
     steps:
       - uses: MathRobin/timezone-action@v1.1

README.md

@@ -1,6 +1,69 @@
-# API to retrieve data from the Belgian IRM KMI in Python
+# Async API to retrieve data from the Belgian IRM KMI in Python
+
+The package exposes the data from the [mobile application of the Belgian IRM KMI](https://www.meteo.be/en/info/faq/products-services/the-rmi-weather-app) as a Python module. 
+
+See more information in the wiki: https://github.com/jdejaegh/irm-kmi-api/wiki
+
+## Quick start example
+
+```python
+import aiohttp
+import asyncio
+from zoneinfo import ZoneInfo
+from irm_kmi_api import IrmKmiApiClientHa
+
+async def print_weather():
+    session = aiohttp.ClientSession()
+    client = IrmKmiApiClientHa(session=session, user_agent="jdejaegh/irm-kmi-api README example")
+    await client.refresh_forecasts_coord({'lat': 50.47, 'long': 4.87})
+    await session.close()
+
+    weather = client.get_current_weather(tz=ZoneInfo('Europe/Brussels'))
+    city = client.get_city()
+
+    print(f"{weather['temperature']}°C with wind of {weather['wind_speed']} km/h in {city}")
+
+
+if __name__ == '__main__':
+    asyncio.run(print_weather())
+```
+
+
+## Features
+
+The package provides the following:  
+1. Current weather
+2. Hourly and daily forecast
+3. Rain radar forecast and animation
+4. Warning data (for extreme weather condition such as storm, thunder, floods)
+5. Pollen data
+
+<details>
+<summary>Screenshots of the rain radar animation</summary>
+<img src="https://github.com/jdejaegh/irm-kmi-api/raw/main/img/camera_light.png"/>  <br>
+<img src="https://github.com/jdejaegh/irm-kmi-api/raw/main/img/camera_dark.png"/>  <br>
+<img src="https://github.com/jdejaegh/irm-kmi-api/raw/main/img/camera_sat.png"/>  
+</details>
+
+## Limitations
+
+The package does not provide the 14-days forcast as in the application.  
+
+This package will not implement any feature that is not available via the API (e.g., humidity and dew point data is not
+provided by the API and thus is not available in this package).
+
+
+## Usage considerations
+
+The API is not publicly documented and has been reversed engineered: it can change at any time without notice and break this package.  
+
+Be mindful when using the API: put a meaningful User-Agent string when creating an `IrmKmiApiClient` and apply rate-limiting for your queries.
+
+
+## Contributions
+
+Contributions are welcome.  Please discuss major changes in an issue before submitting a pull request.
 
-The data is collected via their non-public mobile application API.
 
 
 ## Disclaimer

irm_kmi_api/__init__.py

@@ -1 +1,45 @@
-__version__ = '0.1.6'
+from .api import (
+    IrmKmiApiClient,
+    IrmKmiApiClientHa,
+    IrmKmiApiCommunicationError,
+    IrmKmiApiError,
+)
+from .data import (
+    AnimationFrameData,
+    ConditionEvol,
+    CurrentWeatherData,
+    ExtendedForecast,
+    Forecast,
+    PollenLevel,
+    PollenName,
+    RadarAnimationData,
+    RadarForecast,
+    RadarStyle,
+    WarningData,
+    WarningType,
+)
+from .pollen import PollenParser
+from .rain_graph import RainGraph
+
+__all__ = [
+    "IrmKmiApiClient",
+    "IrmKmiApiClientHa",
+    "IrmKmiApiCommunicationError",
+    "IrmKmiApiError",
+    "AnimationFrameData",
+    "ConditionEvol",
+    "CurrentWeatherData",
+    "ExtendedForecast",
+    "Forecast",
+    "PollenLevel",
+    "PollenName",
+    "RadarAnimationData",
+    "RadarForecast",
+    "RadarStyle",
+    "WarningData",
+    "WarningType",
+    "PollenParser",
+    "RainGraph",
+]
+
+__version__ = '1.1.0'

irm_kmi_api/api.py

@@ -8,18 +8,30 @@ import time
 import urllib.parse
 from datetime import datetime, timedelta
 from statistics import mean
-from typing import List, Dict
+from typing import Dict, List, Tuple
 from zoneinfo import ZoneInfo
 
 import aiohttp
-import async_timeout
 
-from .const import MAP_WARNING_ID_TO_SLUG as SLUG_MAP, WWEVOL_TO_ENUM_MAP
-from .const import STYLE_TO_PARAM_MAP, WEEKDAYS
-from .data import (AnimationFrameData, CurrentWeatherData, Forecast,
-                   IrmKmiForecast, IrmKmiRadarForecast, RadarAnimationData,
-                   WarningData)
-from .pollen import PollenParser
+from .const import (
+    IRM_KMI_TO_HA_CONDITION_MAP,
+    STYLE_TO_PARAM_MAP,
+    WEEKDAYS,
+    WWEVOL_TO_ENUM_MAP,
+)
+from .const import MAP_WARNING_ID_TO_SLUG as SLUG_MAP
+from .data import (
+    AnimationFrameData,
+    CurrentWeatherData,
+    ExtendedForecast,
+    Forecast,
+    RadarAnimationData,
+    RadarForecast,
+    RadarStyle,
+    WarningData,
+    WarningType,
+)
+from .pollen import PollenLevel, PollenName, PollenParser
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,19 +44,21 @@ class IrmKmiApiCommunicationError(IrmKmiApiError):
     """Exception to indicate a communication error."""
 
 
-class IrmKmiApiParametersError(IrmKmiApiError):
-    """Exception to indicate a parameter error."""
-
-
 class IrmKmiApiClient:
     """API client for IRM KMI weather data"""
     COORD_DECIMALS = 6
     _cache_max_age = 60 * 60 * 2  # Remove items from the cache if they have not been hit since 2 hours
     _cache = {}
+    _base_url = "https://app.meteo.be/services/appv4/"
 
     def __init__(self, session: aiohttp.ClientSession, user_agent: str) -> None:
+        """
+        Create a new instance of the API client
+
+        :param session: aiohttp.ClientSession to use for the request
+        :param user_agent: string that will indentify your application in the User-Agent header of the HTTP requests
+        """
         self._session = session
-        self._base_url = "https://app.meteo.be/services/appv4/"
         self._user_agent = user_agent
 
     async def get_forecasts_coord(self, coord: Dict[str, float | int]) -> dict:
@@ -128,7 +142,7 @@ class IrmKmiApiClient:
             headers['If-None-Match'] = self._cache[url]['etag']
 
         try:
-            async with async_timeout.timeout(60):
+            async with asyncio.timeout(60):
                 response = await self._session.request(
                     method=method,
                     url=url,
@@ -167,10 +181,20 @@ class IrmKmiApiClient:
 class IrmKmiApiClientHa(IrmKmiApiClient):
     """API client for IRM KMI weather data with additional methods to integrate easily with Home Assistant"""
 
-    def __init__(self, session: aiohttp.ClientSession, user_agent: str, cdt_map: dict) -> None:
+    def __init__(self, session: aiohttp.ClientSession, user_agent: str, cdt_map: Dict[Tuple[int, str], str] | None = None) -> None:
+        """
+        Create a new instance of the API client.  This client has more methods to integrate easily with Home Assistant
+
+        :param session: aiohttp.ClientSession to use for the request
+        :param user_agent: string that will indentify your application in the User-Agent header of the HTTP requests
+        :param cdt_map: mapping of weather conditions returned by the API and string that should be used when calling the
+                    methods. See the wiki for more information on what conditions are possible:
+                    https://github.com/jdejaegh/irm-kmi-api/wiki/API-documentation#obs-key (Table with icons matching).
+                    Example: cdt_map = { (0, 'd'): 'sunny', (0, 'n'): 'clear_night' }
+        """
         super().__init__(session, user_agent)
         self._api_data = dict()
-        self._cdt_map = cdt_map
+        self._cdt_map = cdt_map if cdt_map is not None else IRM_KMI_TO_HA_CONDITION_MAP
 
     async def refresh_forecasts_coord(self, coord: Dict[str, float | int]) -> None:
         """
@@ -254,7 +278,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
 
         return current_weather
 
-    def get_radar_forecast(self) -> List[IrmKmiRadarForecast]:
+    def get_radar_forecast(self) -> List[RadarForecast]:
         """
         Create a list of short term forecasts for rain based on the data provided by the rain radar
 
@@ -276,7 +300,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
         forecast = list()
         for f in sequence:
             forecast.append(
-                IrmKmiRadarForecast(
+                RadarForecast(
                     datetime=f.get("time"),
                     native_precipitation=f.get('value'),
                     rain_forecast_max=round(f.get('positionHigher') * ratio, 2),
@@ -344,7 +368,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
 
         return forecasts
 
-    def get_daily_forecast(self, tz: ZoneInfo, lang: str) -> List[IrmKmiForecast]:
+    def get_daily_forecast(self, tz: ZoneInfo, lang: str) -> List[ExtendedForecast]:
         """
         Parse the API data we currently have to build the daily forecast list.
 
@@ -416,7 +440,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
                 except (TypeError, ValueError):
                     pass
 
-            forecast = IrmKmiForecast(
+            forecast = ExtendedForecast(
                 datetime=(forecast_day.strftime('%Y-%m-%d')),
                 condition=self._cdt_map.get((f.get('ww1', None), f.get('dayNight', None)), None),
                 condition_2=self._cdt_map.get((f.get('ww2', None), f.get('dayNight', None)), None),
@@ -444,7 +468,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
 
         return forecasts
 
-    def get_animation_data(self, tz: ZoneInfo, lang: str, style: str, dark_mode: bool) -> RadarAnimationData:
+    def get_animation_data(self, tz: ZoneInfo, lang: str, style: RadarStyle, dark_mode: bool) -> RadarAnimationData:
         """
         Get all the image URLs and create the radar animation data object.
 
@@ -525,7 +549,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
 
             result.append(
                 WarningData(
-                    slug=SLUG_MAP.get(warning_id, 'unknown'),
+                    slug=SLUG_MAP.get(warning_id, WarningType.UNKNOWN),
                     id=warning_id,
                     level=level,
                     friendly_name=data.get('warningType', {}).get('name', {}).get(lang, ''),
@@ -537,7 +561,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
 
         return result if len(result) > 0 else []
 
-    async def get_pollen(self) -> Dict[str, str | None]:
+    async def get_pollen(self) -> Dict[PollenName, PollenLevel | None]:
         """
         Get SVG pollen info from the API, return the pollen data dict
 

irm_kmi_api/const.py

@@ -1,39 +1,117 @@
 from typing import Final
 
-from irm_kmi_api.data import IrmKmiConditionEvol
+from .data import ConditionEvol, PollenLevel, RadarStyle, WarningType
+
+POLLEN_LEVEL_TO_COLOR = {
+    'null': PollenLevel.GREEN,
+    'low': PollenLevel.YELLOW,
+    'moderate': PollenLevel.ORANGE,
+    'high': PollenLevel.RED,
+    'very high': PollenLevel.PURPLE,
+    'active': PollenLevel.ACTIVE
+}
 
-POLLEN_NAMES: Final = {'Alder', 'Ash', 'Birch', 'Grasses', 'Hazel', 'Mugwort', 'Oak'}
-POLLEN_LEVEL_TO_COLOR = {'null': 'green', 'low': 'yellow', 'moderate': 'orange', 'high': 'red', 'very high': 'purple',
-                         'active': 'active'}
 WEEKDAYS = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']
 
-# TODO move those to an Enum
-OPTION_STYLE_STD: Final = 'standard_style'
-OPTION_STYLE_CONTRAST: Final = 'contrast_style'
-OPTION_STYLE_YELLOW_RED: Final = 'yellow_red_style'
-OPTION_STYLE_SATELLITE: Final = 'satellite_style'
 STYLE_TO_PARAM_MAP: Final = {
-    OPTION_STYLE_STD: 1,
-    OPTION_STYLE_CONTRAST: 2,
-    OPTION_STYLE_YELLOW_RED: 3,
-    OPTION_STYLE_SATELLITE: 4
+    RadarStyle.OPTION_STYLE_STD: 1,
+    RadarStyle.OPTION_STYLE_CONTRAST: 2,
+    RadarStyle.OPTION_STYLE_YELLOW_RED: 3,
+    RadarStyle.OPTION_STYLE_SATELLITE: 4
 }
 
 MAP_WARNING_ID_TO_SLUG: Final = {
-    0: 'wind',
-    1: 'rain',
-    2: 'ice_or_snow',
-    3: 'thunder',
-    7: 'fog',
-    9: 'cold',
-    12: 'thunder_wind_rain',
-    13: 'thunderstorm_strong_gusts',
-    14: 'thunderstorm_large_rainfall',
-    15: 'storm_surge',
-    17: 'coldspell'}
+    0: WarningType.WIND,
+    1: WarningType.RAIN,
+    2: WarningType.ICE_OR_SNOW,
+    3: WarningType.THUNDER,
+    7: WarningType.FOG,
+    9: WarningType.COLD,
+    10: WarningType.HEAT,
+    12: WarningType.THUNDER_WIND_RAIN,
+    13: WarningType.THUNDERSTORM_STRONG_GUSTS,
+    14: WarningType.THUNDERSTORM_LARGE_RAINFALL,
+    15: WarningType.STORM_SURGE,
+    17: WarningType.COLDSPELL
+}
 
 WWEVOL_TO_ENUM_MAP: Final = {
-    None: IrmKmiConditionEvol.STABLE,
-    0: IrmKmiConditionEvol.ONE_WAY,
-    1: IrmKmiConditionEvol.TWO_WAYS
+    None: ConditionEvol.STABLE,
+    0: ConditionEvol.ONE_WAY,
+    1: ConditionEvol.TWO_WAYS
+}
+
+ATTR_CONDITION_CLEAR_NIGHT = "clear-night"
+ATTR_CONDITION_CLOUDY = "cloudy"
+ATTR_CONDITION_EXCEPTIONAL = "exceptional"
+ATTR_CONDITION_FOG = "fog"
+ATTR_CONDITION_HAIL = "hail"
+ATTR_CONDITION_LIGHTNING = "lightning"
+ATTR_CONDITION_LIGHTNING_RAINY = "lightning-rainy"
+ATTR_CONDITION_PARTLYCLOUDY = "partlycloudy"
+ATTR_CONDITION_POURING = "pouring"
+ATTR_CONDITION_RAINY = "rainy"
+ATTR_CONDITION_SNOWY = "snowy"
+ATTR_CONDITION_SNOWY_RAINY = "snowy-rainy"
+ATTR_CONDITION_SUNNY = "sunny"
+ATTR_CONDITION_WINDY = "windy"
+ATTR_CONDITION_WINDY_VARIANT = "windy-variant"
+
+IRM_KMI_TO_HA_CONDITION_MAP: Final = {
+    (0, 'd'): ATTR_CONDITION_SUNNY,
+    (0, 'n'): ATTR_CONDITION_CLEAR_NIGHT,
+    (1, 'd'): ATTR_CONDITION_SUNNY,
+    (1, 'n'): ATTR_CONDITION_CLEAR_NIGHT,
+    (2, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (2, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (3, 'd'): ATTR_CONDITION_PARTLYCLOUDY,
+    (3, 'n'): ATTR_CONDITION_PARTLYCLOUDY,
+    (4, 'd'): ATTR_CONDITION_POURING,
+    (4, 'n'): ATTR_CONDITION_POURING,
+    (5, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (5, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (6, 'd'): ATTR_CONDITION_POURING,
+    (6, 'n'): ATTR_CONDITION_POURING,
+    (7, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (7, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (8, 'd'): ATTR_CONDITION_SNOWY_RAINY,
+    (8, 'n'): ATTR_CONDITION_SNOWY_RAINY,
+    (9, 'd'): ATTR_CONDITION_SNOWY_RAINY,
+    (9, 'n'): ATTR_CONDITION_SNOWY_RAINY,
+    (10, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (10, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (11, 'd'): ATTR_CONDITION_SNOWY,
+    (11, 'n'): ATTR_CONDITION_SNOWY,
+    (12, 'd'): ATTR_CONDITION_SNOWY,
+    (12, 'n'): ATTR_CONDITION_SNOWY,
+    (13, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (13, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (14, 'd'): ATTR_CONDITION_CLOUDY,
+    (14, 'n'): ATTR_CONDITION_CLOUDY,
+    (15, 'd'): ATTR_CONDITION_CLOUDY,
+    (15, 'n'): ATTR_CONDITION_CLOUDY,
+    (16, 'd'): ATTR_CONDITION_POURING,
+    (16, 'n'): ATTR_CONDITION_POURING,
+    (17, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (17, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
+    (18, 'd'): ATTR_CONDITION_RAINY,
+    (18, 'n'): ATTR_CONDITION_RAINY,
+    (19, 'd'): ATTR_CONDITION_POURING,
+    (19, 'n'): ATTR_CONDITION_POURING,
+    (20, 'd'): ATTR_CONDITION_SNOWY_RAINY,
+    (20, 'n'): ATTR_CONDITION_SNOWY_RAINY,
+    (21, 'd'): ATTR_CONDITION_RAINY,
+    (21, 'n'): ATTR_CONDITION_RAINY,
+    (22, 'd'): ATTR_CONDITION_SNOWY,
+    (22, 'n'): ATTR_CONDITION_SNOWY,
+    (23, 'd'): ATTR_CONDITION_SNOWY,
+    (23, 'n'): ATTR_CONDITION_SNOWY,
+    (24, 'd'): ATTR_CONDITION_FOG,
+    (24, 'n'): ATTR_CONDITION_FOG,
+    (25, 'd'): ATTR_CONDITION_FOG,
+    (25, 'n'): ATTR_CONDITION_FOG,
+    (26, 'd'): ATTR_CONDITION_FOG,
+    (26, 'n'): ATTR_CONDITION_FOG,
+    (27, 'd'): ATTR_CONDITION_FOG,
+    (27, 'n'): ATTR_CONDITION_FOG
 }

irm_kmi_api/data.py

@@ -1,6 +1,6 @@
 """Data classes for IRM KMI integration"""
 from datetime import datetime
-from enum import Enum
+from enum import StrEnum
 from typing import List, Required, TypedDict
 
 
@@ -36,16 +36,66 @@ class Forecast(TypedDict, total=False):
     is_daytime: bool | None  # Mandatory to use with forecast_twice_daily
 
 
-class IrmKmiConditionEvol(Enum):
+class ConditionEvol(StrEnum):
+    """Possible state for evolution between weather conditions"""
+
     ONE_WAY = 'one_way'
     TWO_WAYS = 'two_ways'
     STABLE = 'stable'
 
-class IrmKmiForecast(Forecast, total=False):
+
+class RadarStyle(StrEnum):
+    """Possible style for the rain radar"""
+
+    OPTION_STYLE_STD = 'standard_style'
+    OPTION_STYLE_CONTRAST = 'contrast_style'
+    OPTION_STYLE_YELLOW_RED = 'yellow_red_style'
+    OPTION_STYLE_SATELLITE = 'satellite_style'
+
+
+class PollenName(StrEnum):
+    ALDER = 'alder'
+    ASH = 'ash'
+    BIRCH = 'birch'
+    GRASSES = 'grasses'
+    HAZEL = 'hazel'
+    MUGWORT = 'mugwort'
+    OAK = 'oak'
+
+
+class PollenLevel(StrEnum):
+    """Possible pollen levels"""
+
+    NONE = 'none'
+    ACTIVE = 'active'
+    GREEN = 'green'
+    YELLOW = 'yellow'
+    ORANGE = 'orange'
+    RED = 'red'
+    PURPLE = 'purple'
+
+class WarningType(StrEnum):
+    """Possible warning types"""
+
+    COLD = 'cold'
+    COLDSPELL = 'coldspell'
+    FOG = 'fog'
+    HEAT = 'heat'
+    ICE_OR_SNOW = 'ice_or_snow'
+    RAIN = 'rain'
+    STORM_SURGE = 'storm_surge'
+    THUNDER = 'thunder'
+    THUNDERSTORM_LARGE_RAINFALL = 'thunderstorm_large_rainfall'
+    THUNDERSTORM_STRONG_GUSTS = 'thunderstorm_strong_gusts'
+    THUNDER_WIND_RAIN = 'thunder_wind_rain'
+    WIND = 'wind'
+    UNKNOWN = 'unknown'
+
+class ExtendedForecast(Forecast, total=False):
     """Forecast class with additional attributes for IRM KMI"""
 
     condition_2: str | None
-    condition_evol: IrmKmiConditionEvol | None
+    condition_evol: ConditionEvol | None
     text: str | None
     sunrise: str | None
     sunset: str | None
@@ -53,6 +103,7 @@ class IrmKmiForecast(Forecast, total=False):
 
 class CurrentWeatherData(TypedDict, total=False):
     """Class to hold the currently observable weather at a given location"""
+
     condition: str | None
     temperature: float | None
     wind_speed: float | None
@@ -64,7 +115,8 @@ class CurrentWeatherData(TypedDict, total=False):
 
 class WarningData(TypedDict, total=False):
     """Holds data about a specific warning"""
-    slug: str
+
+    slug: WarningType
     id: int
     level: int
     friendly_name: str
@@ -73,8 +125,9 @@ class WarningData(TypedDict, total=False):
     ends_at: datetime
 
 
-class IrmKmiRadarForecast(Forecast):
+class RadarForecast(Forecast):
     """Forecast class to handle rain forecast from the IRM KMI rain radar"""
+
     rain_forecast_max: float
     rain_forecast_min: float
     might_rain: bool
@@ -83,6 +136,7 @@ class IrmKmiRadarForecast(Forecast):
 
 class AnimationFrameData(TypedDict, total=False):
     """Holds one single frame of the radar camera, along with the timestamp of the frame"""
+
     time: datetime | None
     image: bytes | str | None
     value: float | None
@@ -93,6 +147,7 @@ class AnimationFrameData(TypedDict, total=False):
 
 class RadarAnimationData(TypedDict, total=False):
     """Holds frames and additional data for the animation to be rendered"""
+
     sequence: List[AnimationFrameData] | None
     most_recent_image_idx: int | None
     hint: str | None

irm_kmi_api/pollen.py

@@ -1,9 +1,10 @@
 """Parse pollen info from SVG from IRM KMI api"""
 import logging
 import xml.etree.ElementTree as ET
-from typing import List
+from typing import Dict, List
 
-from .const import POLLEN_LEVEL_TO_COLOR, POLLEN_NAMES
+from .const import POLLEN_LEVEL_TO_COLOR
+from .data import PollenLevel, PollenName
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -21,7 +22,7 @@ class PollenParser:
     ):
         self._xml = xml_string
 
-    def get_pollen_data(self) -> dict:
+    def get_pollen_data(self) -> Dict[PollenName, PollenLevel | None]:
         """
         Parse the SVG and extract the pollen data from the image.
         If an error occurs, return the default value.
@@ -38,11 +39,11 @@ class PollenParser:
 
         elements: List[ET.Element] = self._extract_elements(root)
 
-        pollens = {e.attrib.get('x', None): self._get_elem_text(e).lower()
-                   for e in elements if 'tspan' in e.tag and self._get_elem_text(e) in POLLEN_NAMES}
+        pollens = {e.attrib.get('x', None): self._get_txt(e).lower()
+                   for e in elements if 'tspan' in e.tag and str(self._get_txt(e)).lower() in PollenName}
 
-        pollen_levels = {e.attrib.get('x', None): POLLEN_LEVEL_TO_COLOR[self._get_elem_text(e)]
-                         for e in elements if 'tspan' in e.tag and self._get_elem_text(e) in POLLEN_LEVEL_TO_COLOR}
+        pollen_levels = {e.attrib.get('x', None): POLLEN_LEVEL_TO_COLOR[self._get_txt(e)]
+                         for e in elements if 'tspan' in e.tag and self._get_txt(e) in POLLEN_LEVEL_TO_COLOR}
 
         level_dots = {e.attrib.get('cx', None) for e in elements if 'circle' in e.tag}
 
@@ -50,13 +51,18 @@ class PollenParser:
         # As of January 2025, the text is always 'active' and the dot shows the real level
         # If text says 'active', check the dot; else trust the text
         for position, pollen in pollens.items():
+            # Check if pollen is a known one
+            try:
+                pollen: PollenName = PollenName(pollen)
+            except ValueError:
+                _LOGGER.warning(f'Unknown pollen name {pollen}')
+                continue
             # Determine pollen level based on text
             if position is not None and position in pollen_levels:
                 pollen_data[pollen] = pollen_levels[position]
-                _LOGGER.debug(f"{pollen} is {pollen_data[pollen]} according to text")
-
+                _LOGGER.debug(f"{pollen.value} is {pollen_data[pollen]} according to text")
             # If text is 'active' or if there is no text, check the dot as a fallback
-            if pollen_data[pollen] not in {'none', 'active'}:
+            if pollen_data[pollen] not in {PollenLevel.NONE, PollenLevel.ACTIVE}:
                 _LOGGER.debug(f"{pollen} trusting text")
             else:
                 for dot in level_dots:
@@ -66,35 +72,35 @@ class PollenParser:
                         pass
                     else:
                         if 24 <= relative_x_position <= 34:
-                            pollen_data[pollen] = 'green'
+                            pollen_data[pollen] = PollenLevel.GREEN
                         elif 13 <= relative_x_position <= 23:
-                            pollen_data[pollen] = 'yellow'
+                            pollen_data[pollen] = PollenLevel.YELLOW
                         elif -5 <= relative_x_position <= 5:
-                            pollen_data[pollen] = 'orange'
+                            pollen_data[pollen] = PollenLevel.ORANGE
                         elif -23 <= relative_x_position <= -13:
-                            pollen_data[pollen] = 'red'
+                            pollen_data[pollen] = PollenLevel.RED
                         elif -34 <= relative_x_position <= -24:
-                            pollen_data[pollen] = 'purple'
+                            pollen_data[pollen] = PollenLevel.PURPLE
 
-                _LOGGER.debug(f"{pollen} is {pollen_data[pollen]} according to dot")
+                _LOGGER.debug(f"{pollen.value} is {pollen_data[pollen]} according to dot")
 
         _LOGGER.debug(f"Pollen data: {pollen_data}")
         return pollen_data
 
     @staticmethod
-    def get_default_data() -> dict:
+    def get_default_data() -> Dict[PollenName, PollenLevel | None]:
         """Return all the known pollen with 'none' value"""
-        return {k.lower(): 'none' for k in POLLEN_NAMES}
+        return {k: PollenLevel.NONE for k in PollenName}
 
     @staticmethod
-    def get_unavailable_data() -> dict:
+    def get_unavailable_data() -> Dict[PollenName, PollenLevel | None]:
         """Return all the known pollen with None value"""
-        return {k.lower(): None for k in POLLEN_NAMES}
+        return {k: None for k in PollenName}
 
     @staticmethod
-    def get_option_values() -> List[str]:
+    def get_option_values() -> List[PollenLevel]:
         """List all the values that the pollen can have"""
-        return list(POLLEN_LEVEL_TO_COLOR.values()) + ['none']
+        return list(POLLEN_LEVEL_TO_COLOR.values()) + [PollenLevel.NONE]
 
     @staticmethod
     def _extract_elements(root) -> List[ET.Element]:
@@ -106,7 +112,7 @@ class PollenParser:
         return elements
 
     @staticmethod
-    def _get_elem_text(e) -> str | None:
+    def _get_txt(e) -> str | None:
         if e.text is not None:
             return e.text.strip()
         return None

irm_kmi_api/rain_graph.py

@@ -6,14 +6,12 @@ import datetime
 import logging
 from typing import List, Self
 
-import async_timeout
 from svgwrite import Drawing
 from svgwrite.animate import Animate
 from svgwrite.container import FONT_TEMPLATE
 
 from .api import IrmKmiApiClient, IrmKmiApiError
-from .const import OPTION_STYLE_SATELLITE
-from .data import AnimationFrameData, RadarAnimationData
+from .data import AnimationFrameData, RadarAnimationData, RadarStyle
 from .resources import be_black, be_satellite, be_white, nl, roboto
 
 _LOGGER = logging.getLogger(__name__)
@@ -25,7 +23,7 @@ class RainGraph:
     def __init__(self,
                  animation_data: RadarAnimationData,
                  country: str,
-                 style: str,
+                 style: RadarStyle,
                  dark_mode: bool = False,
                  tz: datetime.tzinfo = None,
                  svg_width: float = 640,
@@ -150,7 +148,7 @@ class RainGraph:
         """
         return self._animation_data.get('hint', '')
 
-    async def _download_clouds(self, idx=None):
+    async def _download_clouds(self, idx: int | None = None):
         """
         Download cloud images and save the result in the internal state.
 
@@ -187,7 +185,7 @@ class RainGraph:
 
         for url in urls:
             coroutines.append(self._api_client.get_image(url))
-        async with async_timeout.timeout(60):
+        async with asyncio.timeout(60):
             images_from_api = await asyncio.gather(*coroutines)
 
         _LOGGER.info(f"Just downloaded {len(images_from_api)} images")
@@ -432,7 +430,7 @@ class RainGraph:
     def _get_background_png_b64(self) -> str:
         if self._country == 'NL':
             return nl.nl_b64
-        elif self._style == OPTION_STYLE_SATELLITE:
+        elif self._style == RadarStyle.OPTION_STYLE_SATELLITE:
             return be_satellite.be_satelitte_b64
         elif self._dark_mode:
             return be_black.be_black_b64

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "irm-kmi-api"
-version = "0.1.6"
+version = "1.1.0"
 description = "Retrieve data from the Belgian IRM KMI in Python"
 readme = "README.md"
 authors = [{ name = "Jules Dejaeghere", email = "curable.grass491@mailer.me" }]
@@ -16,17 +16,19 @@ classifiers = [
 ]
 keywords = ["weather", "weather-api", "netherlands", "weather-forecast", "pollen", "belgium", "luxembourg", "rain-radar"]
 dependencies = [
-    "aiohttp>=3.11.0",
-    "async-timeout>=4.0.3",
-    "svgwrite>=1.4.3",
+    "aiohttp>=3.11.0,<4.0.0",
+    "svgwrite>=1.4.3,<2.0.0",
 ]
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 
 [project.urls]
 Homepage = "https://github.com/jdejaegh/irm-kmi-api"
 
+[tool.setuptools]
+packages = ["irm_kmi_api", "irm_kmi_api.resources"]
+
 [tool.bumpver]
-current_version = "0.1.6"
+current_version = "1.1.0"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "bump version {old_version} -> {new_version}"
 tag_message = "{new_version}"

requirements.txt

@@ -1,3 +1,2 @@
-aiohttp>=3.11.0
-async-timeout>=4.0.3
-svgwrite>=1.4.3
+aiohttp>=3.11.0,<4.0.0
+svgwrite>=1.4.3,<2.0.0

tests/conftest.py

@@ -4,8 +4,8 @@ from __future__ import annotations
 import json
 from unittest.mock import MagicMock
 
-from irm_kmi_api.api import IrmKmiApiClientHa
-from tests.const import IRM_KMI_TO_HA_CONDITION_MAP
+from irm_kmi_api import IrmKmiApiClientHa
+from irm_kmi_api.const import IRM_KMI_TO_HA_CONDITION_MAP
 
 
 def load_fixture(fixture):
@@ -21,3 +21,15 @@ def get_api_with_data(fixture: str) -> IrmKmiApiClientHa:
     api = IrmKmiApiClientHa(session=MagicMock(), user_agent='', cdt_map=IRM_KMI_TO_HA_CONDITION_MAP)
     api._api_data = get_api_data(fixture)
     return api
+
+def is_serializable(x):
+    try:
+        json.dumps(x)
+        return True
+    except (TypeError, OverflowError):
+        return False
+
+def assert_all_serializable(elements: list):
+    for element in elements:
+        for v in element.values():
+            assert is_serializable(v)

tests/const.py

@@ -1,77 +0,0 @@
-from typing import Final
-
-ATTR_CONDITION_CLASS = "condition_class"
-ATTR_CONDITION_CLEAR_NIGHT = "clear-night"
-ATTR_CONDITION_CLOUDY = "cloudy"
-ATTR_CONDITION_EXCEPTIONAL = "exceptional"
-ATTR_CONDITION_FOG = "fog"
-ATTR_CONDITION_HAIL = "hail"
-ATTR_CONDITION_LIGHTNING = "lightning"
-ATTR_CONDITION_LIGHTNING_RAINY = "lightning-rainy"
-ATTR_CONDITION_PARTLYCLOUDY = "partlycloudy"
-ATTR_CONDITION_POURING = "pouring"
-ATTR_CONDITION_RAINY = "rainy"
-ATTR_CONDITION_SNOWY = "snowy"
-ATTR_CONDITION_SNOWY_RAINY = "snowy-rainy"
-ATTR_CONDITION_SUNNY = "sunny"
-ATTR_CONDITION_WINDY = "windy"
-ATTR_CONDITION_WINDY_VARIANT = "windy-variant"
-
-IRM_KMI_TO_HA_CONDITION_MAP: Final = {
-    (0, 'd'): ATTR_CONDITION_SUNNY,
-    (0, 'n'): ATTR_CONDITION_CLEAR_NIGHT,
-    (1, 'd'): ATTR_CONDITION_SUNNY,
-    (1, 'n'): ATTR_CONDITION_CLEAR_NIGHT,
-    (2, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (2, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (3, 'd'): ATTR_CONDITION_PARTLYCLOUDY,
-    (3, 'n'): ATTR_CONDITION_PARTLYCLOUDY,
-    (4, 'd'): ATTR_CONDITION_POURING,
-    (4, 'n'): ATTR_CONDITION_POURING,
-    (5, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (5, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (6, 'd'): ATTR_CONDITION_POURING,
-    (6, 'n'): ATTR_CONDITION_POURING,
-    (7, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (7, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (8, 'd'): ATTR_CONDITION_SNOWY_RAINY,
-    (8, 'n'): ATTR_CONDITION_SNOWY_RAINY,
-    (9, 'd'): ATTR_CONDITION_SNOWY_RAINY,
-    (9, 'n'): ATTR_CONDITION_SNOWY_RAINY,
-    (10, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (10, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (11, 'd'): ATTR_CONDITION_SNOWY,
-    (11, 'n'): ATTR_CONDITION_SNOWY,
-    (12, 'd'): ATTR_CONDITION_SNOWY,
-    (12, 'n'): ATTR_CONDITION_SNOWY,
-    (13, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (13, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (14, 'd'): ATTR_CONDITION_CLOUDY,
-    (14, 'n'): ATTR_CONDITION_CLOUDY,
-    (15, 'd'): ATTR_CONDITION_CLOUDY,
-    (15, 'n'): ATTR_CONDITION_CLOUDY,
-    (16, 'd'): ATTR_CONDITION_POURING,
-    (16, 'n'): ATTR_CONDITION_POURING,
-    (17, 'd'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (17, 'n'): ATTR_CONDITION_LIGHTNING_RAINY,
-    (18, 'd'): ATTR_CONDITION_RAINY,
-    (18, 'n'): ATTR_CONDITION_RAINY,
-    (19, 'd'): ATTR_CONDITION_POURING,
-    (19, 'n'): ATTR_CONDITION_POURING,
-    (20, 'd'): ATTR_CONDITION_SNOWY_RAINY,
-    (20, 'n'): ATTR_CONDITION_SNOWY_RAINY,
-    (21, 'd'): ATTR_CONDITION_RAINY,
-    (21, 'n'): ATTR_CONDITION_RAINY,
-    (22, 'd'): ATTR_CONDITION_SNOWY,
-    (22, 'n'): ATTR_CONDITION_SNOWY,
-    (23, 'd'): ATTR_CONDITION_SNOWY,
-    (23, 'n'): ATTR_CONDITION_SNOWY,
-    (24, 'd'): ATTR_CONDITION_FOG,
-    (24, 'n'): ATTR_CONDITION_FOG,
-    (25, 'd'): ATTR_CONDITION_FOG,
-    (25, 'n'): ATTR_CONDITION_FOG,
-    (26, 'd'): ATTR_CONDITION_FOG,
-    (26, 'n'): ATTR_CONDITION_FOG,
-    (27, 'd'): ATTR_CONDITION_FOG,
-    (27, 'n'): ATTR_CONDITION_FOG
-}

tests/test_api.py

@@ -1,15 +1,16 @@
 import json
 import time
-from datetime import datetime as dt, timedelta
+from datetime import datetime as dt
+from datetime import timedelta
 from unittest.mock import AsyncMock, MagicMock
 from zoneinfo import ZoneInfo
 
 import freezegun
 import pytest
 
-from irm_kmi_api.api import IrmKmiApiClient, IrmKmiApiClientHa
-from irm_kmi_api.data import CurrentWeatherData
-from irm_kmi_api.pollen import PollenParser
+from irm_kmi_api import IrmKmiApiClient, IrmKmiApiClientHa
+from irm_kmi_api import CurrentWeatherData
+from irm_kmi_api import PollenParser
 
 
 @freezegun.freeze_time(dt.fromisoformat('2025-05-03T17:30:00+00:00'))

tests/test_current_weather.py

@@ -4,9 +4,9 @@ from zoneinfo import ZoneInfo
 import pytest
 from freezegun import freeze_time
 
-from irm_kmi_api.data import CurrentWeatherData
-from tests.conftest import get_api_data, get_api_with_data
-from tests.const import ATTR_CONDITION_CLOUDY, ATTR_CONDITION_PARTLYCLOUDY
+from irm_kmi_api import CurrentWeatherData
+from tests.conftest import get_api_data, get_api_with_data, is_serializable
+from irm_kmi_api.const import ATTR_CONDITION_CLOUDY, ATTR_CONDITION_PARTLYCLOUDY
 
 
 @freeze_time(datetime.fromisoformat('2023-12-26T17:30:00+00:00'))
@@ -140,3 +140,10 @@ def test_current_weather_attributes(
         assert r == expected_
 
     run(sensor, expected)
+
+@freeze_time(datetime.fromisoformat('2023-12-26T17:30:00+00:00'))
+def test_current_weather_is_serializable() -> None:
+    api = get_api_with_data("forecast.json")
+    tz = ZoneInfo("Europe/Brussels")
+    result = api.get_current_weather(tz)
+    assert is_serializable(result)

tests/test_daily_forecast.py

@@ -3,9 +3,9 @@ from zoneinfo import ZoneInfo
 
 from freezegun import freeze_time
 
-from irm_kmi_api.data import IrmKmiForecast, IrmKmiConditionEvol
-from tests.conftest import get_api_with_data
-from tests.const import ATTR_CONDITION_PARTLYCLOUDY
+from irm_kmi_api import ConditionEvol, ExtendedForecast
+from tests.conftest import get_api_with_data, assert_all_serializable
+from irm_kmi_api.const import ATTR_CONDITION_PARTLYCLOUDY
 
 
 @freeze_time(datetime.fromisoformat('2023-12-26T18:30:00+01:00'))
@@ -19,11 +19,11 @@ async def test_daily_forecast() -> None:
     assert len(result) == 8
     assert result[0]['datetime'] == '2023-12-26'
     assert not result[0]['is_daytime']
-    expected = IrmKmiForecast(
+    expected = ExtendedForecast(
         datetime='2023-12-27',
         condition=ATTR_CONDITION_PARTLYCLOUDY,
         condition_2=None,
-        condition_evol=IrmKmiConditionEvol.TWO_WAYS,
+        condition_evol=ConditionEvol.TWO_WAYS,
         native_precipitation=0,
         native_temperature=9,
         native_templow=4,
@@ -106,3 +106,11 @@ async def test_sunrise_sunset_be() -> None:
 
     assert result[2]['sunrise'] == '2023-12-28T08:45:00+01:00'
     assert result[2]['sunset'] == '2023-12-28T16:43:00+01:00'
+
+
+def test_daily_serializable() -> None:
+    api = get_api_with_data("forecast.json")
+    tz = ZoneInfo("Europe/Brussels")
+
+    result = api.get_daily_forecast(tz, 'fr')
+    assert_all_serializable(result)

tests/test_hourly_forecast.py

@@ -3,9 +3,9 @@ from zoneinfo import ZoneInfo
 
 from freezegun import freeze_time
 
-from irm_kmi_api.data import Forecast
-from tests.conftest import get_api_with_data
-from tests.const import ATTR_CONDITION_CLOUDY, ATTR_CONDITION_RAINY
+from irm_kmi_api import Forecast
+from tests.conftest import get_api_with_data, assert_all_serializable
+from irm_kmi_api.const import ATTR_CONDITION_CLOUDY, ATTR_CONDITION_RAINY
 
 
 @freeze_time(datetime.fromisoformat('2023-12-26T18:30:00+01:00'))
@@ -88,4 +88,10 @@ def test_hourly_forecast_midnight_bug() -> None:
 
     assert result[24]['datetime'] == '2024-06-01T00:00:00+02:00'
 
+def test_hourly_serializable() -> None:
+    api = get_api_with_data("forecast.json")
+    tz = ZoneInfo("Europe/Brussels")
 
+    result = api.get_hourly_forecast(tz)
+
+    assert_all_serializable(result)

tests/test_pollen.py

@@ -1,37 +1,63 @@
 from unittest.mock import AsyncMock
 
-from irm_kmi_api.pollen import PollenParser
-from tests.conftest import get_api_with_data, load_fixture
+from irm_kmi_api import PollenLevel, PollenName, PollenParser
+from tests.conftest import get_api_with_data, load_fixture, is_serializable
 
 
 def test_svg_pollen_parsing():
     with open("tests/fixtures/pollen.svg", "r") as file:
         svg_data = file.read()
     data = PollenParser(svg_data).get_pollen_data()
-    assert data == {'birch': 'none', 'oak': 'none', 'hazel': 'none', 'mugwort': 'none', 'alder': 'none',
-                    'grasses': 'purple', 'ash': 'none'}
+    assert data == {PollenName.BIRCH: PollenLevel.NONE,
+                    PollenName.OAK: PollenLevel.NONE,
+                    PollenName.HAZEL: PollenLevel.NONE,
+                    PollenName.MUGWORT: PollenLevel.NONE,
+                    PollenName.ALDER: PollenLevel.NONE,
+                    PollenName.GRASSES: PollenLevel.PURPLE,
+                    PollenName.ASH: PollenLevel.NONE}
 
 def test_svg_two_pollen_parsing():
     with open("tests/fixtures/new_two_pollens.svg", "r") as file:
         svg_data = file.read()
     data = PollenParser(svg_data).get_pollen_data()
-    assert data == {'birch': 'none', 'oak': 'none', 'hazel': 'none', 'mugwort': 'active', 'alder': 'none',
-                    'grasses': 'red', 'ash': 'none'}
+    assert data == {PollenName.BIRCH: PollenLevel.NONE,
+                    PollenName.OAK: PollenLevel.NONE,
+                    PollenName.HAZEL: PollenLevel.NONE,
+                    PollenName.MUGWORT: PollenLevel.ACTIVE,
+                    PollenName.ALDER: PollenLevel.NONE,
+                    PollenName.GRASSES: PollenLevel.RED,
+                    PollenName.ASH: PollenLevel.NONE}
 
 def test_svg_two_pollen_parsing_2025_update():
     with open("tests/fixtures/pollens-2025.svg", "r") as file:
         svg_data = file.read()
     data = PollenParser(svg_data).get_pollen_data()
-    assert data == {'birch': 'none', 'oak': 'none', 'hazel': 'active', 'mugwort': 'none', 'alder': 'green',
-                    'grasses': 'none', 'ash': 'none'}
+    assert data == {PollenName.BIRCH: PollenLevel.NONE,
+                    PollenName.OAK: PollenLevel.NONE,
+                    PollenName.HAZEL: PollenLevel.ACTIVE,
+                    PollenName.MUGWORT: PollenLevel.NONE,
+                    PollenName.ALDER: PollenLevel.GREEN,
+                    PollenName.GRASSES: PollenLevel.NONE,
+                    PollenName.ASH: PollenLevel.NONE}
 
 def test_pollen_options():
-    assert set(PollenParser.get_option_values()) == {'green', 'yellow', 'orange', 'red', 'purple', 'active', 'none'}
+    assert set(PollenParser.get_option_values()) == {PollenLevel.GREEN,
+                                                     PollenLevel.YELLOW,
+                                                     PollenLevel.ORANGE,
+                                                     PollenLevel.RED,
+                                                     PollenLevel.PURPLE,
+                                                     PollenLevel.ACTIVE,
+                                                     PollenLevel.NONE}
 
 
 def test_pollen_default_values():
-    assert PollenParser.get_default_data() == {'birch': 'none', 'oak': 'none', 'hazel': 'none', 'mugwort': 'none',
-                                               'alder': 'none', 'grasses': 'none', 'ash': 'none'}
+    assert PollenParser.get_default_data() == {PollenName.BIRCH: PollenLevel.NONE,
+                                               PollenName.OAK: PollenLevel.NONE,
+                                               PollenName.HAZEL: PollenLevel.NONE,
+                                               PollenName.MUGWORT: PollenLevel.NONE,
+                                               PollenName.ALDER: PollenLevel.NONE,
+                                               PollenName.GRASSES: PollenLevel.NONE,
+                                               PollenName.ASH: PollenLevel.NONE}
 
 
 async def test_pollen_data_from_api() -> None:
@@ -41,7 +67,18 @@ async def test_pollen_data_from_api() -> None:
     api.get_svg = AsyncMock(return_value=load_fixture("pollen.svg"))
 
     result = await api.get_pollen()
-    expected = {'mugwort': 'none', 'birch': 'none', 'alder': 'none', 'ash': 'none', 'oak': 'none',
-                'grasses': 'purple', 'hazel': 'none'}
+    expected = {PollenName.MUGWORT: PollenLevel.NONE,
+                PollenName.BIRCH: PollenLevel.NONE,
+                PollenName.ALDER: PollenLevel.NONE,
+                PollenName.ASH: PollenLevel.NONE,
+                PollenName.OAK: PollenLevel.NONE,
+                PollenName.GRASSES: PollenLevel.PURPLE,
+                PollenName.HAZEL: PollenLevel.NONE}
     assert result == expected
 
+def test_pollen_is_serializable():
+    with open("tests/fixtures/pollens-2025.svg", "r") as file:
+        svg_data = file.read()
+    data = PollenParser(svg_data).get_pollen_data()
+
+    assert is_serializable(data)

tests/test_radar_forecast.py

@@ -1,7 +1,7 @@
 import pytest
 
-from irm_kmi_api.data import IrmKmiRadarForecast
-from tests.conftest import get_api_with_data
+from irm_kmi_api import RadarForecast
+from tests.conftest import get_api_with_data, assert_all_serializable
 
 
 def test_radar_forecast() -> None:
@@ -9,28 +9,28 @@ def test_radar_forecast() -> None:
     result = api.get_radar_forecast()
 
     expected = [
-        IrmKmiRadarForecast(datetime="2023-12-26T17:00:00+01:00", native_precipitation=0, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T17:10:00+01:00", native_precipitation=0, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T17:20:00+01:00", native_precipitation=0, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T17:30:00+01:00", native_precipitation=0, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T17:40:00+01:00", native_precipitation=0.1, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T17:50:00+01:00", native_precipitation=0.01, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T18:00:00+01:00", native_precipitation=0.12, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T18:10:00+01:00", native_precipitation=1.2, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T18:20:00+01:00", native_precipitation=2, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T18:30:00+01:00", native_precipitation=0, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
-        IrmKmiRadarForecast(datetime="2023-12-26T18:40:00+01:00", native_precipitation=0, might_rain=False,
-                            rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min')
+        RadarForecast(datetime="2023-12-26T17:00:00+01:00", native_precipitation=0, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T17:10:00+01:00", native_precipitation=0, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T17:20:00+01:00", native_precipitation=0, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T17:30:00+01:00", native_precipitation=0, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T17:40:00+01:00", native_precipitation=0.1, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T17:50:00+01:00", native_precipitation=0.01, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T18:00:00+01:00", native_precipitation=0.12, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T18:10:00+01:00", native_precipitation=1.2, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T18:20:00+01:00", native_precipitation=2, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T18:30:00+01:00", native_precipitation=0, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min'),
+        RadarForecast(datetime="2023-12-26T18:40:00+01:00", native_precipitation=0, might_rain=False,
+                      rain_forecast_max=0, rain_forecast_min=0, unit='mm/10min')
     ]
 
     assert expected == result
@@ -40,7 +40,7 @@ def test_radar_forecast_rain_interval() -> None:
     api = get_api_with_data('forecast_with_rain_on_radar.json')
     result = api.get_radar_forecast()
 
-    _12 = IrmKmiRadarForecast(
+    _12 = RadarForecast(
         datetime='2024-05-30T18:00:00+02:00',
         native_precipitation=0.89,
         might_rain=True,
@@ -49,7 +49,7 @@ def test_radar_forecast_rain_interval() -> None:
         unit='mm/10min'
     )
 
-    _13 = IrmKmiRadarForecast(
+    _13 = RadarForecast(
         datetime="2024-05-30T18:10:00+02:00",
         native_precipitation=0.83,
         might_rain=True,
@@ -78,3 +78,10 @@ async def test_current_rainfall_unit(
 
     for r in radar_forecast:
         assert r.get('unit') == expected
+
+def test_radar_serializable() -> None:
+    api = get_api_with_data("forecast.json")
+
+    result = api.get_radar_forecast()
+
+    assert_all_serializable(result)

tests/test_rain_graph.py

@@ -3,13 +3,12 @@ import datetime
 import json
 from datetime import datetime as dt
 from datetime import timedelta
-from unittest.mock import MagicMock, AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 from zoneinfo import ZoneInfo
 
-from irm_kmi_api.api import IrmKmiApiClientHa
-from irm_kmi_api.const import OPTION_STYLE_SATELLITE
-from irm_kmi_api.data import AnimationFrameData, RadarAnimationData
-from irm_kmi_api.rain_graph import RainGraph
+from irm_kmi_api import IrmKmiApiClientHa
+from irm_kmi_api import AnimationFrameData, RadarAnimationData, RadarStyle
+from irm_kmi_api import RainGraph
 from tests.conftest import load_fixture
 
 
@@ -45,7 +44,7 @@ async def test_svg_frame_setup():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     await rain_graph._draw_svg_frame()
@@ -65,7 +64,7 @@ def test_svg_hint():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._write_hint()
@@ -81,7 +80,7 @@ def test_svg_time_bars():
         tz = datetime.UTC,
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._draw_hour_bars()
@@ -100,7 +99,7 @@ def test_draw_chances_path():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._draw_chances_path()
@@ -118,7 +117,7 @@ def test_draw_data_line():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._draw_data_line()
@@ -136,7 +135,7 @@ async def test_insert_background():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     await rain_graph._insert_background()
@@ -159,7 +158,7 @@ def test_draw_current_frame_line_moving():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._draw_current_fame_line()
@@ -187,7 +186,7 @@ def test_draw_current_frame_line_index():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._draw_current_fame_line(0)
@@ -216,7 +215,7 @@ def test_draw_description_text():
         tz=datetime.UTC,
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._draw_description_text()
@@ -243,7 +242,7 @@ def test_draw_cloud_layer():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._insert_cloud_layer()
@@ -263,7 +262,7 @@ async def test_draw_location_layer():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     await rain_graph._draw_location()
@@ -281,7 +280,7 @@ def test_get_animation_data():
 
     tz = ZoneInfo('Europe/Brussels')
     lang = 'en'
-    style = OPTION_STYLE_SATELLITE
+    style = RadarStyle.OPTION_STYLE_SATELLITE
     dark_mode = False
 
     api._api_data = json.loads(load_fixture("forecast.json"))
@@ -306,7 +305,7 @@ async def test_download_single_cloud():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._api_client = MagicMock()
@@ -324,7 +323,7 @@ async def test_download_many_clouds():
     rain_graph = RainGraph(
         animation_data=data,
         country='BE',
-        style='STD',
+        style=RadarStyle.OPTION_STYLE_STD,
     )
 
     rain_graph._api_client = MagicMock()
@@ -339,11 +338,11 @@ def test_can_build_rain_graph_with_empty_sequence():
 
     RainGraph(
         RadarAnimationData(sequence=None),
-        'en', 'style'
+        'BE', RadarStyle.OPTION_STYLE_STD
     )
 
     RainGraph(
         RadarAnimationData(sequence=[]),
-        'en', 'style'
+        'BE', RadarStyle.OPTION_STYLE_STD
     )
 

tests/test_warning.py

@@ -2,7 +2,8 @@ from datetime import datetime
 
 from freezegun import freeze_time
 
-from tests.conftest import get_api_with_data
+from irm_kmi_api import WarningType
+from tests.conftest import get_api_with_data, is_serializable
 
 
 @freeze_time(datetime.fromisoformat('2024-01-12T07:10:00+00:00'))
@@ -18,7 +19,32 @@ async def test_warning_data() -> None:
     assert first.get('starts_at').replace(tzinfo=None) < datetime.now()
     assert first.get('ends_at').replace(tzinfo=None) > datetime.now()
 
-    assert first.get('slug') == 'fog'
+    assert first.get('slug') == WarningType.FOG
     assert first.get('friendly_name') == 'Fog'
     assert first.get('id') == 7
     assert first.get('level') == 1
+
+async def test_warning_heat() -> None:
+    api = get_api_with_data("antwerp_with_heat_warning.json")
+
+    result = api.get_warnings(lang='en')
+
+    assert isinstance(result, list)
+    assert len(result) == 1
+
+    first = result[0]
+
+    assert first.get('slug') == WarningType.HEAT
+    assert first.get('friendly_name') == 'Heat'
+    assert first.get('id') == 10
+    assert first.get('level') == 1
+
+
+async def test_warning_data_is_serializable() -> None:
+    api = get_api_with_data("be_forecast_warning.json")
+
+    result = api.get_warnings(lang='en')
+    for r in result:
+        del r["starts_at"]
+        del r["ends_at"]
+        assert is_serializable(r)