diff --git a/irm_kmi_api/api.py b/irm_kmi_api/api.py index bb818c2..32a1e3e 100644 --- a/irm_kmi_api/api.py +++ b/irm_kmi_api/api.py @@ -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 .data import (AnimationFrameData, CurrentWeatherData, Forecast, ExtendedForecast, IrmKmiRadarForecast, RadarAnimationData, - WarningData, RadarStyle) + WarningData, RadarStyle, WarningType) from .pollen import PollenParser _LOGGER = logging.getLogger(__name__) @@ -525,7 +525,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, ''), diff --git a/irm_kmi_api/const.py b/irm_kmi_api/const.py index e6025cd..5082348 100644 --- a/irm_kmi_api/const.py +++ b/irm_kmi_api/const.py @@ -1,14 +1,14 @@ from typing import Final -from .data import ConditionEvol, RadarStyle, PollenLevels +from .data import ConditionEvol, RadarStyle, PollenLevel, WarningType POLLEN_LEVEL_TO_COLOR = { - 'null': PollenLevels.GREEN, - 'low': PollenLevels.YELLOW, - 'moderate': PollenLevels.ORANGE, - 'high': PollenLevels.RED, - 'very high': PollenLevels.PURPLE, - 'active': PollenLevels.ACTIVE + 'null': PollenLevel.GREEN, + 'low': PollenLevel.YELLOW, + 'moderate': PollenLevel.ORANGE, + 'high': PollenLevel.RED, + 'very high': PollenLevel.PURPLE, + 'active': PollenLevel.ACTIVE } WEEKDAYS = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday'] @@ -21,20 +21,21 @@ STYLE_TO_PARAM_MAP: Final = { } 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, + 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: ConditionEvol.STABLE, 0: ConditionEvol.ONE_WAY, 1: ConditionEvol.TWO_WAYS -} \ No newline at end of file +} diff --git a/irm_kmi_api/data.py b/irm_kmi_api/data.py index 2bd45f3..5a03feb 100644 --- a/irm_kmi_api/data.py +++ b/irm_kmi_api/data.py @@ -53,7 +53,7 @@ class RadarStyle(Enum): OPTION_STYLE_SATELLITE = 'satellite_style' -class PollenNames(Enum): +class PollenName(Enum): """Pollens names from the API""" ALDER = 'alder' @@ -65,7 +65,7 @@ class PollenNames(Enum): OAK = 'oak' -class PollenLevels(Enum): +class PollenLevel(Enum): """Possible pollen levels""" NONE = 'none' @@ -76,6 +76,21 @@ class PollenLevels(Enum): RED = 'red' 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): """Forecast class with additional attributes for IRM KMI""" @@ -102,7 +117,7 @@ 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 diff --git a/irm_kmi_api/pollen.py b/irm_kmi_api/pollen.py index 1883e76..f4a209e 100644 --- a/irm_kmi_api/pollen.py +++ b/irm_kmi_api/pollen.py @@ -4,7 +4,7 @@ import xml.etree.ElementTree as ET from typing import List, Dict from .const import POLLEN_LEVEL_TO_COLOR -from .data import PollenNames, PollenLevels +from .data import PollenName, PollenLevel _LOGGER = logging.getLogger(__name__) @@ -22,7 +22,7 @@ class PollenParser: ): 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. If an error occurs, return the default value. @@ -40,7 +40,7 @@ class PollenParser: elements: List[ET.Element] = self._extract_elements(root) 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)] 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(): # Check if pollen is a known one try: - pollen: PollenNames = PollenNames(pollen) + pollen: PollenName = PollenName(pollen) except ValueError: _LOGGER.warning(f'Unknown pollen name {pollen}') continue @@ -62,7 +62,7 @@ class PollenParser: pollen_data[pollen] = pollen_levels[position] _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 {PollenLevels.NONE, PollenLevels.ACTIVE}: + if pollen_data[pollen] not in {PollenLevel.NONE, PollenLevel.ACTIVE}: _LOGGER.debug(f"{pollen} trusting text") else: for dot in level_dots: @@ -72,15 +72,15 @@ class PollenParser: pass else: if 24 <= relative_x_position <= 34: - pollen_data[pollen] = PollenLevels.GREEN + pollen_data[pollen] = PollenLevel.GREEN elif 13 <= relative_x_position <= 23: - pollen_data[pollen] = PollenLevels.YELLOW + pollen_data[pollen] = PollenLevel.YELLOW elif -5 <= relative_x_position <= 5: - pollen_data[pollen] = PollenLevels.ORANGE + pollen_data[pollen] = PollenLevel.ORANGE elif -23 <= relative_x_position <= -13: - pollen_data[pollen] = PollenLevels.RED + pollen_data[pollen] = PollenLevel.RED 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") @@ -88,19 +88,19 @@ class PollenParser: return pollen_data @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 {k: PollenLevels.NONE for k in PollenNames} + return {k: PollenLevel.NONE for k in PollenName} @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 {k: None for k in PollenNames} + return {k: None for k in PollenName} @staticmethod - def get_option_values() -> List[PollenLevels]: + def get_option_values() -> List[PollenLevel]: """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 def _extract_elements(root) -> List[ET.Element]: diff --git a/tests/test_pollen.py b/tests/test_pollen.py index 8549582..ce9751b 100644 --- a/tests/test_pollen.py +++ b/tests/test_pollen.py @@ -1,7 +1,7 @@ import logging 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 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: svg_data = file.read() data = PollenParser(svg_data).get_pollen_data() - assert data == {PollenNames.BIRCH: PollenLevels.NONE, - PollenNames.OAK: PollenLevels.NONE, - PollenNames.HAZEL: PollenLevels.NONE, - PollenNames.MUGWORT: PollenLevels.NONE, - PollenNames.ALDER: PollenLevels.NONE, - PollenNames.GRASSES: PollenLevels.PURPLE, - PollenNames.ASH: PollenLevels.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 == {PollenNames.BIRCH: PollenLevels.NONE, - PollenNames.OAK: PollenLevels.NONE, - PollenNames.HAZEL: PollenLevels.NONE, - PollenNames.MUGWORT: PollenLevels.ACTIVE, - PollenNames.ALDER: PollenLevels.NONE, - PollenNames.GRASSES: PollenLevels.RED, - PollenNames.ASH: PollenLevels.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 == {PollenNames.BIRCH: PollenLevels.NONE, - PollenNames.OAK: PollenLevels.NONE, - PollenNames.HAZEL: PollenLevels.ACTIVE, - PollenNames.MUGWORT: PollenLevels.NONE, - PollenNames.ALDER: PollenLevels.GREEN, - PollenNames.GRASSES: PollenLevels.NONE, - PollenNames.ASH: PollenLevels.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()) == {PollenLevels.GREEN, - PollenLevels.YELLOW, - PollenLevels.ORANGE, - PollenLevels.RED, - PollenLevels.PURPLE, - PollenLevels.ACTIVE, - PollenLevels.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() == {PollenNames.BIRCH: PollenLevels.NONE, - PollenNames.OAK: PollenLevels.NONE, - PollenNames.HAZEL: PollenLevels.NONE, - PollenNames.MUGWORT: PollenLevels.NONE, - PollenNames.ALDER: PollenLevels.NONE, - PollenNames.GRASSES: PollenLevels.NONE, - PollenNames.ASH: PollenLevels.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: @@ -69,12 +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 = {PollenNames.MUGWORT: PollenLevels.NONE, - PollenNames.BIRCH: PollenLevels.NONE, - PollenNames.ALDER: PollenLevels.NONE, - PollenNames.ASH: PollenLevels.NONE, - PollenNames.OAK: PollenLevels.NONE, - PollenNames.GRASSES: PollenLevels.PURPLE, - PollenNames.HAZEL: PollenLevels.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 diff --git a/tests/test_warning.py b/tests/test_warning.py index 1da98c6..8f16310 100644 --- a/tests/test_warning.py +++ b/tests/test_warning.py @@ -2,6 +2,7 @@ from datetime import datetime from freezegun import freeze_time +from irm_kmi_api.data import WarningType 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('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