jdejaegh/irm-kmi-api
1
0
Fork 0
mirror of https://github.com/jdejaegh/irm-kmi-api.git synced 2025-06-26 20:05:40 +02:00
Code Issues Projects Releases Packages Wiki Activity

Breaking: use Enum for pollen data

Browse source
This commit is contained in:
Jules 2025-05-05 20:11:10 +02:00
parent c4433f20cc
commit af40cba92d
Signed by: jdejaegh
GPG key ID: 99D6D184CA66933A
4 changed files with 95 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
irm_kmi_api/const.py
View file

@ -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']

26
irm_kmi_api/data.py
View file

@ -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"""

45
irm_kmi_api/pollen.py
View file

@ -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]:

55
tests/test_pollen.py
View file

@ -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