mirror of
https://github.com/jdejaegh/irm-kmi-api.git
synced 2025-06-27 04:05:56 +02:00
Use more enums
This commit is contained in:
parent
c06f1c8972
commit
c728493542
GPG key ID:
99D6D184CA66933A
6 changed files with 101 additions and 84 deletions
|
|
@ -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, ''),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue