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

Use more enums

Browse source
This commit is contained in:
Jules 2025-05-05 22:17:54 +02:00
parent c06f1c8972
commit c728493542
Signed by: jdejaegh
GPG key ID: 99D6D184CA66933A
6 changed files with 101 additions and 84 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
irm_kmi_api/api.py
View file

@ -18,7 +18,7 @@ from .const import MAP_WARNING_ID_TO_SLUG as SLUG_MAP, WWEVOL_TO_ENUM_MAP
from .const import STYLE_TO_PARAM_MAP, WEEKDAYS from .const import STYLE_TO_PARAM_MAP, WEEKDAYS
from .data import (AnimationFrameData, CurrentWeatherData, Forecast, from .data import (AnimationFrameData, CurrentWeatherData, Forecast,
ExtendedForecast, IrmKmiRadarForecast, RadarAnimationData, ExtendedForecast, IrmKmiRadarForecast, RadarAnimationData,
WarningData, RadarStyle) WarningData, RadarStyle, WarningType)
from .pollen import PollenParser from .pollen import PollenParser
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -525,7 +525,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
result.append( result.append(
WarningData( WarningData(
slug=SLUG_MAP.get(warning_id, 'unknown'), slug=SLUG_MAP.get(warning_id, WarningType.UNKNOWN),
id=warning_id, id=warning_id,
level=level, level=level,
friendly_name=data.get('warningType', {}).get('name', {}).get(lang, ''), friendly_name=data.get('warningType', {}).get('name', {}).get(lang, ''),

37
irm_kmi_api/const.py
View file

@ -1,14 +1,14 @@
from typing import Final from typing import Final
from .data import ConditionEvol, RadarStyle, PollenLevels from .data import ConditionEvol, RadarStyle, PollenLevel, WarningType
POLLEN_LEVEL_TO_COLOR = { POLLEN_LEVEL_TO_COLOR = {
'null': PollenLevels.GREEN, 'null': PollenLevel.GREEN,
'low': PollenLevels.YELLOW, 'low': PollenLevel.YELLOW,
'moderate': PollenLevels.ORANGE, 'moderate': PollenLevel.ORANGE,
'high': PollenLevels.RED, 'high': PollenLevel.RED,
'very high': PollenLevels.PURPLE, 'very high': PollenLevel.PURPLE,
'active': PollenLevels.ACTIVE 'active': PollenLevel.ACTIVE
} }
WEEKDAYS = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday'] WEEKDAYS = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']
@ -21,17 +21,18 @@ STYLE_TO_PARAM_MAP: Final = {
} }
MAP_WARNING_ID_TO_SLUG: Final = { MAP_WARNING_ID_TO_SLUG: Final = {
0: 'wind', 0: WarningType.WIND,
1: 'rain', 1: WarningType.RAIN,
2: 'ice_or_snow', 2: WarningType.ICE_OR_SNOW,
3: 'thunder', 3: WarningType.THUNDER,
7: 'fog', 7: WarningType.FOG,
9: 'cold', 9: WarningType.COLD,
12: 'thunder_wind_rain', 12: WarningType.THUNDER_WIND_RAIN,
13: 'thunderstorm_strong_gusts', 13: WarningType.THUNDERSTORM_STRONG_GUSTS,
14: 'thunderstorm_large_rainfall', 14: WarningType.THUNDERSTORM_LARGE_RAINFALL,
15: 'storm_surge', 15: WarningType.STORM_SURGE,
17: 'coldspell'} 17: WarningType.COLDSPELL
}
WWEVOL_TO_ENUM_MAP: Final = { WWEVOL_TO_ENUM_MAP: Final = {
None: ConditionEvol.STABLE, None: ConditionEvol.STABLE,

21
irm_kmi_api/data.py
View file

@ -53,7 +53,7 @@ class RadarStyle(Enum):
OPTION_STYLE_SATELLITE = 'satellite_style' OPTION_STYLE_SATELLITE = 'satellite_style'
class PollenNames(Enum): class PollenName(Enum):
"""Pollens names from the API""" """Pollens names from the API"""
ALDER = 'alder' ALDER = 'alder'
@ -65,7 +65,7 @@ class PollenNames(Enum):
OAK = 'oak' OAK = 'oak'
class PollenLevels(Enum): class PollenLevel(Enum):
"""Possible pollen levels""" """Possible pollen levels"""
NONE = 'none' NONE = 'none'
@ -76,6 +76,21 @@ class PollenLevels(Enum):
RED = 'red' RED = 'red'
PURPLE = 'purple' PURPLE = 'purple'
class WarningType(Enum):
"""Possible warning types"""
COLD = 'cold'
COLDSPELL = 'coldspell'
FOG = 'fog'
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): class ExtendedForecast(Forecast, total=False):
"""Forecast class with additional attributes for IRM KMI""" """Forecast class with additional attributes for IRM KMI"""
@ -102,7 +117,7 @@ class CurrentWeatherData(TypedDict, total=False):
class WarningData(TypedDict, total=False): class WarningData(TypedDict, total=False):
"""Holds data about a specific warning""" """Holds data about a specific warning"""
slug: str slug: WarningType
id: int id: int
level: int level: int
friendly_name: str friendly_name: str

32
irm_kmi_api/pollen.py
View file

@ -4,7 +4,7 @@ import xml.etree.ElementTree as ET
from typing import List, Dict from typing import List, Dict
from .const import POLLEN_LEVEL_TO_COLOR from .const import POLLEN_LEVEL_TO_COLOR
from .data import PollenNames, PollenLevels from .data import PollenName, PollenLevel
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -22,7 +22,7 @@ class PollenParser:
): ):
self._xml = xml_string self._xml = xml_string
def get_pollen_data(self) -> Dict[PollenNames, PollenLevels | None]: def get_pollen_data(self) -> Dict[PollenName, PollenLevel | None]:
""" """
Parse the SVG and extract the pollen data from the image. Parse the SVG and extract the pollen data from the image.
If an error occurs, return the default value. If an error occurs, return the default value.
@ -40,7 +40,7 @@ class PollenParser:
elements: List[ET.Element] = self._extract_elements(root) elements: List[ET.Element] = self._extract_elements(root)
pollens = {e.attrib.get('x', None): self._get_txt(e).lower() 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 PollenNames} 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_txt(e)] 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} for e in elements if 'tspan' in e.tag and self._get_txt(e) in POLLEN_LEVEL_TO_COLOR}
@ -53,7 +53,7 @@ class PollenParser:
for position, pollen in pollens.items(): for position, pollen in pollens.items():
# Check if pollen is a known one # Check if pollen is a known one
try: try:
pollen: PollenNames = PollenNames(pollen) pollen: PollenName = PollenName(pollen)
except ValueError: except ValueError:
_LOGGER.warning(f'Unknown pollen name {pollen}') _LOGGER.warning(f'Unknown pollen name {pollen}')
continue continue
@ -62,7 +62,7 @@ class PollenParser:
pollen_data[pollen] = pollen_levels[position] pollen_data[pollen] = pollen_levels[position]
_LOGGER.debug(f"{pollen.value} 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 text is 'active' or if there is no text, check the dot as a fallback
if pollen_data[pollen] not in {PollenLevels.NONE, PollenLevels.ACTIVE}: if pollen_data[pollen] not in {PollenLevel.NONE, PollenLevel.ACTIVE}:
_LOGGER.debug(f"{pollen} trusting text") _LOGGER.debug(f"{pollen} trusting text")
else: else:
for dot in level_dots: for dot in level_dots:
@ -72,15 +72,15 @@ class PollenParser:
pass pass
else: else:
if 24 <= relative_x_position <= 34: if 24 <= relative_x_position <= 34:
pollen_data[pollen] = PollenLevels.GREEN pollen_data[pollen] = PollenLevel.GREEN
elif 13 <= relative_x_position <= 23: elif 13 <= relative_x_position <= 23:
pollen_data[pollen] = PollenLevels.YELLOW pollen_data[pollen] = PollenLevel.YELLOW
elif -5 <= relative_x_position <= 5: elif -5 <= relative_x_position <= 5:
pollen_data[pollen] = PollenLevels.ORANGE pollen_data[pollen] = PollenLevel.ORANGE
elif -23 <= relative_x_position <= -13: elif -23 <= relative_x_position <= -13:
pollen_data[pollen] = PollenLevels.RED pollen_data[pollen] = PollenLevel.RED
elif -34 <= relative_x_position <= -24: elif -34 <= relative_x_position <= -24:
pollen_data[pollen] = PollenLevels.PURPLE pollen_data[pollen] = PollenLevel.PURPLE
_LOGGER.debug(f"{pollen.value} is {pollen_data[pollen]} according to dot") _LOGGER.debug(f"{pollen.value} is {pollen_data[pollen]} according to dot")
@ -88,19 +88,19 @@ class PollenParser:
return pollen_data return pollen_data
@staticmethod @staticmethod
def get_default_data() -> Dict[PollenNames, PollenLevels | None]: def get_default_data() -> Dict[PollenName, PollenLevel | None]:
"""Return all the known pollen with 'none' value""" """Return all the known pollen with 'none' value"""
return {k: PollenLevels.NONE for k in PollenNames} return {k: PollenLevel.NONE for k in PollenName}
@staticmethod @staticmethod
def get_unavailable_data() -> Dict[PollenNames, PollenLevels | None]: def get_unavailable_data() -> Dict[PollenName, PollenLevel | None]:
"""Return all the known pollen with None value""" """Return all the known pollen with None value"""
return {k: None for k in PollenNames} return {k: None for k in PollenName}
@staticmethod @staticmethod
def get_option_values() -> List[PollenLevels]: def get_option_values() -> List[PollenLevel]:
"""List all the values that the pollen can have""" """List all the values that the pollen can have"""
return list(POLLEN_LEVEL_TO_COLOR.values()) + [PollenLevels.NONE] return list(POLLEN_LEVEL_TO_COLOR.values()) + [PollenLevel.NONE]
@staticmethod @staticmethod
def _extract_elements(root) -> List[ET.Element]: def _extract_elements(root) -> List[ET.Element]:

86
tests/test_pollen.py
View file

@ -1,7 +1,7 @@
import logging import logging
from unittest.mock import AsyncMock from unittest.mock import AsyncMock
from irm_kmi_api.data import PollenNames, PollenLevels from irm_kmi_api.data import PollenName, PollenLevel
from irm_kmi_api.pollen import PollenParser from irm_kmi_api.pollen import PollenParser
from tests.conftest import get_api_with_data, load_fixture from tests.conftest import get_api_with_data, load_fixture
@ -10,56 +10,56 @@ def test_svg_pollen_parsing():
with open("tests/fixtures/pollen.svg", "r") as file: with open("tests/fixtures/pollen.svg", "r") as file:
svg_data = file.read() svg_data = file.read()
data = PollenParser(svg_data).get_pollen_data() data = PollenParser(svg_data).get_pollen_data()
assert data == {PollenNames.BIRCH: PollenLevels.NONE, assert data == {PollenName.BIRCH: PollenLevel.NONE,
PollenNames.OAK: PollenLevels.NONE, PollenName.OAK: PollenLevel.NONE,
PollenNames.HAZEL: PollenLevels.NONE, PollenName.HAZEL: PollenLevel.NONE,
PollenNames.MUGWORT: PollenLevels.NONE, PollenName.MUGWORT: PollenLevel.NONE,
PollenNames.ALDER: PollenLevels.NONE, PollenName.ALDER: PollenLevel.NONE,
PollenNames.GRASSES: PollenLevels.PURPLE, PollenName.GRASSES: PollenLevel.PURPLE,
PollenNames.ASH: PollenLevels.NONE} PollenName.ASH: PollenLevel.NONE}
def test_svg_two_pollen_parsing(): def test_svg_two_pollen_parsing():
with open("tests/fixtures/new_two_pollens.svg", "r") as file: with open("tests/fixtures/new_two_pollens.svg", "r") as file:
svg_data = file.read() svg_data = file.read()
data = PollenParser(svg_data).get_pollen_data() data = PollenParser(svg_data).get_pollen_data()
assert data == {PollenNames.BIRCH: PollenLevels.NONE, assert data == {PollenName.BIRCH: PollenLevel.NONE,
PollenNames.OAK: PollenLevels.NONE, PollenName.OAK: PollenLevel.NONE,
PollenNames.HAZEL: PollenLevels.NONE, PollenName.HAZEL: PollenLevel.NONE,
PollenNames.MUGWORT: PollenLevels.ACTIVE, PollenName.MUGWORT: PollenLevel.ACTIVE,
PollenNames.ALDER: PollenLevels.NONE, PollenName.ALDER: PollenLevel.NONE,
PollenNames.GRASSES: PollenLevels.RED, PollenName.GRASSES: PollenLevel.RED,
PollenNames.ASH: PollenLevels.NONE} PollenName.ASH: PollenLevel.NONE}
def test_svg_two_pollen_parsing_2025_update(): def test_svg_two_pollen_parsing_2025_update():
with open("tests/fixtures/pollens-2025.svg", "r") as file: with open("tests/fixtures/pollens-2025.svg", "r") as file:
svg_data = file.read() svg_data = file.read()
data = PollenParser(svg_data).get_pollen_data() data = PollenParser(svg_data).get_pollen_data()
assert data == {PollenNames.BIRCH: PollenLevels.NONE, assert data == {PollenName.BIRCH: PollenLevel.NONE,
PollenNames.OAK: PollenLevels.NONE, PollenName.OAK: PollenLevel.NONE,
PollenNames.HAZEL: PollenLevels.ACTIVE, PollenName.HAZEL: PollenLevel.ACTIVE,
PollenNames.MUGWORT: PollenLevels.NONE, PollenName.MUGWORT: PollenLevel.NONE,
PollenNames.ALDER: PollenLevels.GREEN, PollenName.ALDER: PollenLevel.GREEN,
PollenNames.GRASSES: PollenLevels.NONE, PollenName.GRASSES: PollenLevel.NONE,
PollenNames.ASH: PollenLevels.NONE} PollenName.ASH: PollenLevel.NONE}
def test_pollen_options(): def test_pollen_options():
assert set(PollenParser.get_option_values()) == {PollenLevels.GREEN, assert set(PollenParser.get_option_values()) == {PollenLevel.GREEN,
PollenLevels.YELLOW, PollenLevel.YELLOW,
PollenLevels.ORANGE, PollenLevel.ORANGE,
PollenLevels.RED, PollenLevel.RED,
PollenLevels.PURPLE, PollenLevel.PURPLE,
PollenLevels.ACTIVE, PollenLevel.ACTIVE,
PollenLevels.NONE} PollenLevel.NONE}
def test_pollen_default_values(): def test_pollen_default_values():
assert PollenParser.get_default_data() == {PollenNames.BIRCH: PollenLevels.NONE, assert PollenParser.get_default_data() == {PollenName.BIRCH: PollenLevel.NONE,
PollenNames.OAK: PollenLevels.NONE, PollenName.OAK: PollenLevel.NONE,
PollenNames.HAZEL: PollenLevels.NONE, PollenName.HAZEL: PollenLevel.NONE,
PollenNames.MUGWORT: PollenLevels.NONE, PollenName.MUGWORT: PollenLevel.NONE,
PollenNames.ALDER: PollenLevels.NONE, PollenName.ALDER: PollenLevel.NONE,
PollenNames.GRASSES: PollenLevels.NONE, PollenName.GRASSES: PollenLevel.NONE,
PollenNames.ASH: PollenLevels.NONE} PollenName.ASH: PollenLevel.NONE}
async def test_pollen_data_from_api() -> None: async def test_pollen_data_from_api() -> None:
@ -69,12 +69,12 @@ async def test_pollen_data_from_api() -> None:
api.get_svg = AsyncMock(return_value=load_fixture("pollen.svg")) api.get_svg = AsyncMock(return_value=load_fixture("pollen.svg"))
result = await api.get_pollen() result = await api.get_pollen()
expected = {PollenNames.MUGWORT: PollenLevels.NONE, expected = {PollenName.MUGWORT: PollenLevel.NONE,
PollenNames.BIRCH: PollenLevels.NONE, PollenName.BIRCH: PollenLevel.NONE,
PollenNames.ALDER: PollenLevels.NONE, PollenName.ALDER: PollenLevel.NONE,
PollenNames.ASH: PollenLevels.NONE, PollenName.ASH: PollenLevel.NONE,
PollenNames.OAK: PollenLevels.NONE, PollenName.OAK: PollenLevel.NONE,
PollenNames.GRASSES: PollenLevels.PURPLE, PollenName.GRASSES: PollenLevel.PURPLE,
PollenNames.HAZEL: PollenLevels.NONE} PollenName.HAZEL: PollenLevel.NONE}
assert result == expected assert result == expected

3
tests/test_warning.py
View file

@ -2,6 +2,7 @@ from datetime import datetime
from freezegun import freeze_time from freezegun import freeze_time
from irm_kmi_api.data import WarningType
from tests.conftest import get_api_with_data from tests.conftest import get_api_with_data
@ -18,7 +19,7 @@ async def test_warning_data() -> None:
assert first.get('starts_at').replace(tzinfo=None) < datetime.now() assert first.get('starts_at').replace(tzinfo=None) < datetime.now()
assert first.get('ends_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('friendly_name') == 'Fog'
assert first.get('id') == 7 assert first.get('id') == 7
assert first.get('level') == 1 assert first.get('level') == 1