Use more enums

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

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 .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, ''),

View file

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

View file

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

View file

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

View file

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

View file

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