diff --git a/irm_kmi_api/const.py b/irm_kmi_api/const.py index 6e47e09..d5bfde1 100644 --- a/irm_kmi_api/const.py +++ b/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', - 'low': 'yellow', - 'moderate': 'orange', - 'high': 'red', - 'very high': 'purple', - 'active': 'active' + 'null': IrmKmiPollenLevels.GREEN, + 'low': IrmKmiPollenLevels.YELLOW, + 'moderate': IrmKmiPollenLevels.ORANGE, + 'high': IrmKmiPollenLevels.RED, + 'very high': IrmKmiPollenLevels.PURPLE, + 'active': IrmKmiPollenLevels.ACTIVE } WEEKDAYS = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday'] diff --git a/irm_kmi_api/data.py b/irm_kmi_api/data.py index 9116602..0d2f3ac 100644 --- a/irm_kmi_api/data.py +++ b/irm_kmi_api/data.py @@ -52,13 +52,25 @@ class IrmKmiRadarStyle(Enum): class IrmKmiPollenNames(Enum): """Pollens names from the API""" - ALDER = 'Alder' - ASH = 'Ash' - BIRCH = 'Birch' - GRASSES = 'Grasses' - HAZEL = 'Hazel' - MUGWORT = 'Mugwort' - OAK = 'Oak' + ALDER = 'alder' + ASH = 'ash' + BIRCH = 'birch' + GRASSES = 'grasses' + HAZEL = 'hazel' + MUGWORT = 'mugwort' + 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""" diff --git a/irm_kmi_api/pollen.py b/irm_kmi_api/pollen.py index f736342..9c5a9a4 100644 --- a/irm_kmi_api/pollen.py +++ b/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]: diff --git a/tests/test_pollen.py b/tests/test_pollen.py index 64051b8..f63bb29 100644 --- a/tests/test_pollen.py +++ b/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', - 'grasses': 'purple', 'ash': 'none'} + assert data == {IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.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', - 'grasses': 'red', 'ash': 'none'} + assert data == {IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.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', - 'grasses': 'none', 'ash': 'none'} + assert data == {IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.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', - 'alder': 'none', 'grasses': 'none', 'ash': 'none'} + assert PollenParser.get_default_data() == {IrmKmiPollenNames.BIRCH: IrmKmiPollenLevels.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', - 'grasses': 'purple', 'hazel': 'none'} + expected = {IrmKmiPollenNames.MUGWORT: IrmKmiPollenLevels.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