Breaking: use Enum for pollen data

irm_kmi_api/const.py

@@ -1,14 +1,14 @@
 from typing import Final
 
-from .data import IrmKmiConditionEvol, IrmKmiRadarStyle
+from .data import IrmKmiConditionEvol, IrmKmiRadarStyle, IrmKmiPollenLevels
 
 POLLEN_LEVEL_TO_COLOR = {
-    'null': 'green',
+    'null': IrmKmiPollenLevels.GREEN,
-    'low': 'yellow',
+    'low': IrmKmiPollenLevels.YELLOW,
-    'moderate': 'orange',
+    'moderate': IrmKmiPollenLevels.ORANGE,
-    'high': 'red',
+    'high': IrmKmiPollenLevels.RED,
-    'very high': 'purple',
+    'very high': IrmKmiPollenLevels.PURPLE,
-    'active': 'active'
+    'active': IrmKmiPollenLevels.ACTIVE
 }
 
 WEEKDAYS = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']

irm_kmi_api/data.py

@@ -52,13 +52,25 @@ class IrmKmiRadarStyle(Enum):
 class IrmKmiPollenNames(Enum):
     """Pollens names from the API"""
 
-    ALDER = 'Alder'
+    ALDER = 'alder'
-    ASH = 'Ash'
+    ASH = 'ash'
-    BIRCH = 'Birch'
+    BIRCH = 'birch'
-    GRASSES = 'Grasses'
+    GRASSES = 'grasses'
-    HAZEL = 'Hazel'
+    HAZEL = 'hazel'
-    MUGWORT = 'Mugwort'
+    MUGWORT = 'mugwort'
-    OAK = 'Oak'
+    OAK = 'oak'
+
+
+class IrmKmiPollenLevels(Enum):
+    """Possible pollen levels"""
+
+    NONE = 'none'
+    ACTIVE = 'active'
+    GREEN = 'green'
+    YELLOW = 'yellow'
+    ORANGE = 'orange'
+    RED = 'red'
+    PURPLE = 'purple'
 
 class IrmKmiForecast(Forecast, total=False):
     """Forecast class with additional attributes for IRM KMI"""

irm_kmi_api/pollen.py

@@ -1,10 +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 List, Dict
 
 from .const import POLLEN_LEVEL_TO_COLOR
-from .data import IrmKmiPollenNames
+from .data import IrmKmiPollenNames, IrmKmiPollenLevels
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -22,7 +22,7 @@ class PollenParser:
     ):
         self._xml = xml_string
 
-    def get_pollen_data(self) -> dict:
+    def get_pollen_data(self) -> Dict[IrmKmiPollenNames, IrmKmiPollenLevels | None]:
         """
         Parse the SVG and extract the pollen data from the image.
         If an error occurs, return the default value.
@@ -40,10 +40,10 @@ 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 IrmKmiPollenNames}
+                   for e in elements if 'tspan' in e.tag and str(self._get_elem_text(e)).lower() in IrmKmiPollenNames}
 
         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 IrmKmiPollenNames}
+                         for e in elements if 'tspan' in e.tag and self._get_elem_text(e) in POLLEN_LEVEL_TO_COLOR}
 
         level_dots = {e.attrib.get('cx', None) for e in elements if 'circle' in e.tag}
 
@@ -51,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: IrmKmiPollenNames = IrmKmiPollenNames(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 {IrmKmiPollenLevels.NONE, IrmKmiPollenLevels.ACTIVE}:
                 _LOGGER.debug(f"{pollen} trusting text")
             else:
                 for dot in level_dots:
@@ -67,35 +72,35 @@ class PollenParser:
                         pass
                     else:
                         if 24 <= relative_x_position <= 34:
-                            pollen_data[pollen] = 'green'
+                            pollen_data[pollen] = IrmKmiPollenLevels.GREEN
                         elif 13 <= relative_x_position <= 23:
-                            pollen_data[pollen] = 'yellow'
+                            pollen_data[pollen] = IrmKmiPollenLevels.YELLOW
                         elif -5 <= relative_x_position <= 5:
-                            pollen_data[pollen] = 'orange'
+                            pollen_data[pollen] = IrmKmiPollenLevels.ORANGE
                         elif -23 <= relative_x_position <= -13:
-                            pollen_data[pollen] = 'red'
+                            pollen_data[pollen] = IrmKmiPollenLevels.RED
                         elif -34 <= relative_x_position <= -24:
-                            pollen_data[pollen] = 'purple'
+                            pollen_data[pollen] = IrmKmiPollenLevels.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[IrmKmiPollenNames, IrmKmiPollenLevels | None]:
         """Return all the known pollen with 'none' value"""
-        return {k.value.lower(): 'none' for k in IrmKmiPollenNames}
+        return {k: IrmKmiPollenLevels.NONE for k in IrmKmiPollenNames}
 
     @staticmethod
-    def get_unavailable_data() -> dict:
+    def get_unavailable_data() -> Dict[IrmKmiPollenNames, IrmKmiPollenLevels | None]:
         """Return all the known pollen with None value"""
-        return {k.value.lower(): None for k in IrmKmiPollenNames}
+        return {k: None for k in IrmKmiPollenNames}
 
     @staticmethod
-    def get_option_values() -> List[str]:
+    def get_option_values() -> List[IrmKmiPollenLevels]:
         """List all the values that the pollen can have"""
-        return list(POLLEN_LEVEL_TO_COLOR.values()) + ['none']
+        return list(POLLEN_LEVEL_TO_COLOR.values()) + [IrmKmiPollenLevels.NONE]
 
     @staticmethod
     def _extract_elements(root) -> List[ET.Element]:

tests/test_pollen.py

@@ -1,5 +1,7 @@
+import logging
 from unittest.mock import AsyncMock
 
+from irm_kmi_api.data import IrmKmiPollenNames, IrmKmiPollenLevels
 from irm_kmi_api.pollen import PollenParser
 from tests.conftest import get_api_with_data, load_fixture
 
@@ -8,30 +10,56 @@ 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',
+    assert data == {IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.NONE, 
-                    'grasses': 'purple', 'ash': 'none'}
+                    IrmKmiPollenNames.OAK: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.HAZEL: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.MUGWORT: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.ALDER: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.GRASSES: IrmKmiPollenLevels.PURPLE, 
+                    IrmKmiPollenNames.ASH: IrmKmiPollenLevels.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',
+    assert data == {IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.NONE,
-                    'grasses': 'red', 'ash': 'none'}
+                    IrmKmiPollenNames.OAK: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.HAZEL: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.MUGWORT: IrmKmiPollenLevels.ACTIVE,
+                    IrmKmiPollenNames.ALDER: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.GRASSES: IrmKmiPollenLevels.RED,
+                    IrmKmiPollenNames.ASH: IrmKmiPollenLevels.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',
+    assert data == {IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.NONE,
-                    'grasses': 'none', 'ash': 'none'}
+                    IrmKmiPollenNames.OAK: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.HAZEL: IrmKmiPollenLevels.ACTIVE,
+                    IrmKmiPollenNames.MUGWORT: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.ALDER: IrmKmiPollenLevels.GREEN,
+                    IrmKmiPollenNames.GRASSES: IrmKmiPollenLevels.NONE,
+                    IrmKmiPollenNames.ASH: IrmKmiPollenLevels.NONE}
 
 def test_pollen_options():
-    assert set(PollenParser.get_option_values()) == {'green', 'yellow', 'orange', 'red', 'purple', 'active', 'none'}
+    assert set(PollenParser.get_option_values()) == {IrmKmiPollenLevels.GREEN,
+                                                     IrmKmiPollenLevels.YELLOW,
+                                                     IrmKmiPollenLevels.ORANGE,
+                                                     IrmKmiPollenLevels.RED,
+                                                     IrmKmiPollenLevels.PURPLE,
+                                                     IrmKmiPollenLevels.ACTIVE,
+                                                     IrmKmiPollenLevels.NONE}
 
 
 def test_pollen_default_values():
-    assert PollenParser.get_default_data() == {'birch': 'none', 'oak': 'none', 'hazel': 'none', 'mugwort': 'none',
+    assert PollenParser.get_default_data() == {IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.NONE,
-                                               'alder': 'none', 'grasses': 'none', 'ash': 'none'}
+                                               IrmKmiPollenNames.OAK: IrmKmiPollenLevels.NONE,
+                                               IrmKmiPollenNames.HAZEL: IrmKmiPollenLevels.NONE,
+                                               IrmKmiPollenNames.MUGWORT: IrmKmiPollenLevels.NONE,
+                                               IrmKmiPollenNames.ALDER: IrmKmiPollenLevels.NONE,
+                                               IrmKmiPollenNames.GRASSES: IrmKmiPollenLevels.NONE,
+                                               IrmKmiPollenNames.ASH: IrmKmiPollenLevels.NONE}
 
 
 async def test_pollen_data_from_api() -> None:
@@ -41,7 +69,12 @@ 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',
+    expected = {IrmKmiPollenNames.MUGWORT: IrmKmiPollenLevels.NONE,
-                'grasses': 'purple', 'hazel': 'none'}
+                IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.NONE,
+                IrmKmiPollenNames.ALDER: IrmKmiPollenLevels.NONE,
+                IrmKmiPollenNames.ASH: IrmKmiPollenLevels.NONE,
+                IrmKmiPollenNames.OAK: IrmKmiPollenLevels.NONE,
+                IrmKmiPollenNames.GRASSES: IrmKmiPollenLevels.PURPLE,
+                IrmKmiPollenNames.HAZEL: IrmKmiPollenLevels.NONE}
     assert result == expected