diff --git a/custom_components/irm_kmi/const.py b/custom_components/irm_kmi/const.py
index d1fa951..0ff7f76 100644
--- a/custom_components/irm_kmi/const.py
+++ b/custom_components/irm_kmi/const.py
@@ -143,6 +143,7 @@
17: 'coldspell'}
POLLEN_NAMES: Final = {'Alder', 'Ash', 'Birch', 'Grasses', 'Hazel', 'Mugwort', 'Oak'}
+POLLEN_LEVEL_TO_COLOR = {'null': 'green', 'low': 'yellow', 'moderate': 'orange', 'high': 'red', 'very high': 'purple'}
POLLEN_TO_ICON_MAP: Final = {
'alder': 'mdi:tree', 'ash': 'mdi:tree', 'birch': 'mdi:tree', 'grasses': 'mdi:grass', 'hazel': 'mdi:tree',
diff --git a/custom_components/irm_kmi/pollen.py b/custom_components/irm_kmi/pollen.py
index e29c729..b402abe 100644
--- a/custom_components/irm_kmi/pollen.py
+++ b/custom_components/irm_kmi/pollen.py
@@ -3,32 +3,16 @@
import xml.etree.ElementTree as ET
from typing import List
-from custom_components.irm_kmi.const import POLLEN_NAMES
+from custom_components.irm_kmi.const import POLLEN_LEVEL_TO_COLOR, POLLEN_NAMES
_LOGGER = logging.getLogger(__name__)
-def get_unavailable_data() -> dict:
- """Return all the known pollen with 'none' value"""
- return {k.lower(): 'none' for k in POLLEN_NAMES}
-
-
class PollenParser:
"""
- The SVG looks as follows (see test fixture for a real example)
-
- Active pollens
- ---------------------------------
- Oak active
- Ash active
- ---------------------------------
- Birch ---|---|---|---|-*-
- Alder -*-|---|---|---|---
-
- This classe parses the oak and ash as active, birch as purple and alder as green in the example.
- For active pollen, check if an active text is present on the same line as the pollen name
- For the color scale, look for a white dot (nearly) at the same level as the pollen name. From the white dot
- horizontal position, determine the level
+ Extract pollen level from an SVG provided by the IRM KMI API.
+ To get the data, match pollen names and pollen levels that are vertically aligned. Then, map the value to the
+ corresponding color on the scale.
"""
def __init__(
@@ -37,23 +21,6 @@ def __init__(
):
self._xml = xml_string
- @staticmethod
- def _validate_svg(elements: List[ET.Element]) -> bool:
- """Make sure that the colors of the scale are still where we expect them"""
- x_values = {"rectgreen": 80,
- "rectyellow": 95,
- "rectorange": 110,
- "rectred": 125,
- "rectpurple": 140}
- for e in elements:
- if e.attrib.get('id', '') in x_values.keys():
- try:
- if float(e.attrib.get('x', '0')) != x_values.get(e.attrib.get('id')):
- return False
- except ValueError:
- return False
- return True
-
@staticmethod
def get_default_data() -> dict:
"""Return all the known pollen with 'none' value"""
@@ -67,7 +34,7 @@ def get_unavailable_data() -> dict:
@staticmethod
def get_option_values() -> List[str]:
"""List all the values that the pollen can have"""
- return ['active', 'green', 'yellow', 'orange', 'red', 'purple', 'none']
+ return list(POLLEN_LEVEL_TO_COLOR.values()) + ['none']
@staticmethod
def _extract_elements(root) -> List[ET.Element]:
@@ -79,27 +46,10 @@ def _extract_elements(root) -> List[ET.Element]:
return elements
@staticmethod
- def _dot_to_color_value(dot: ET.Element) -> str:
- """Map the dot horizontal position to a color or 'none'"""
- try:
- cx = float(dot.attrib.get('cx'))
- except ValueError:
- return 'none'
-
- if cx > 155:
- return 'none'
- elif cx > 140:
- return 'purple'
- elif cx > 125:
- return 'red'
- elif cx > 110:
- return 'orange'
- elif cx > 95:
- return 'yellow'
- elif cx > 80:
- return 'green'
- else:
- return 'none'
+ def _get_elem_text(e) -> str | None:
+ if e.text is not None:
+ return e.text.strip()
+ return None
def get_pollen_data(self) -> dict:
"""From the XML string, parse the SVG and extract the pollen data from the image.
@@ -114,28 +64,15 @@ def get_pollen_data(self) -> dict:
elements: List[ET.Element] = self._extract_elements(root)
- if not self._validate_svg(elements):
- _LOGGER.warning("Could not validate SVG pollen data")
- return pollen_data
+ 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 POLLEN_NAMES}
+
+ 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 POLLEN_LEVEL_TO_COLOR}
- pollens = [e for e in elements if 'tspan' in e.tag and e.text in POLLEN_NAMES]
- active = [e for e in elements if 'tspan' in e.tag and e.text == 'active']
- dots = [e for e in elements if 'ellipse' in e.tag
- and 'fill:#ffffff' in e.attrib.get('style', '')
- and 3 == float(e.attrib.get('rx', '0'))]
-
- for pollen in pollens:
- try:
- y = float(pollen.attrib.get('y'))
- if y in [float(e.attrib.get('y')) for e in active]:
- pollen_data[pollen.text.lower()] = 'active'
- else:
- dot = [d for d in dots if y - 3 <= float(d.attrib.get('cy', '0')) <= y + 3]
- if len(dot) == 1:
- dot = dot[0]
- pollen_data[pollen.text.lower()] = self._dot_to_color_value(dot)
- except ValueError | NameError:
- _LOGGER.warning("Skipped some data in the pollen SVG")
+ for position, pollen in pollens.items():
+ if position is not None and position in pollen_levels:
+ pollen_data[pollen] = pollen_levels[position]
_LOGGER.debug(f"Pollen data: {pollen_data}")
return pollen_data
diff --git a/custom_components/irm_kmi/sensor.py b/custom_components/irm_kmi/sensor.py
index 64aaf0f..baeaca3 100644
--- a/custom_components/irm_kmi/sensor.py
+++ b/custom_components/irm_kmi/sensor.py
@@ -1,6 +1,6 @@
"""Sensor for pollen from the IRM KMI"""
-from datetime import datetime
import logging
+from datetime import datetime
from homeassistant.components import sensor
from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
diff --git a/tests/fixtures/pollen.svg b/tests/fixtures/pollen.svg
index 0c23b2a..affd362 100644
--- a/tests/fixtures/pollen.svg
+++ b/tests/fixtures/pollen.svg
@@ -1,2 +1,42 @@
-
\ No newline at end of file
+
\ No newline at end of file
diff --git a/tests/fixtures/pollen_three.svg b/tests/fixtures/pollen_three.svg
deleted file mode 100644
index 04db508..0000000
--- a/tests/fixtures/pollen_three.svg
+++ /dev/null
@@ -1,45 +0,0 @@
-
-
\ No newline at end of file
diff --git a/tests/fixtures/pollen_two.svg b/tests/fixtures/pollen_two.svg
deleted file mode 100644
index 4bfe1b6..0000000
--- a/tests/fixtures/pollen_two.svg
+++ /dev/null
@@ -1,2 +0,0 @@
-
-
\ No newline at end of file
diff --git a/tests/test_pollen.py b/tests/test_pollen.py
index b868e8c..f358cb3 100644
--- a/tests/test_pollen.py
+++ b/tests/test_pollen.py
@@ -12,24 +12,12 @@ 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': 'purple', 'oak': 'active', 'hazel': 'none', 'mugwort': 'none', 'alder': 'green',
- 'grasses': 'none', 'ash': 'active'}
-
- with open("tests/fixtures/pollen_two.svg", "r") as file:
- svg_data = file.read()
- data = PollenParser(svg_data).get_pollen_data()
- assert data == {'birch': 'purple', 'oak': 'active', 'hazel': 'none', 'mugwort': 'none', 'alder': 'active',
- 'grasses': 'none', 'ash': 'active'}
-
- with open("tests/fixtures/pollen_three.svg", "r") as file:
- svg_data = file.read()
- data = PollenParser(svg_data).get_pollen_data()
- assert data == {'birch': 'none', 'oak': 'active', 'hazel': 'none', 'mugwort': 'none', 'alder': 'active',
- 'grasses': 'none', 'ash': 'active'}
+ assert data == {'birch': 'none', 'oak': 'none', 'hazel': 'none', 'mugwort': 'none', 'alder': 'none',
+ 'grasses': 'purple', 'ash': 'none'}
def test_pollen_options():
- assert PollenParser.get_option_values() == ['active', 'green', 'yellow', 'orange', 'red', 'purple', 'none']
+ assert set(PollenParser.get_option_values()) == {'green', 'yellow', 'orange', 'red', 'purple', 'none'}
def test_pollen_default_values():
@@ -46,8 +34,8 @@ async def test_pollen_data_from_api(
api_data = get_api_data("be_forecast_warning.json")
result = await coordinator._async_pollen_data(api_data)
- expected = {'mugwort': 'none', 'birch': 'purple', 'alder': 'green', 'ash': 'active', 'oak': 'active',
- 'grasses': 'none', 'hazel': 'none'}
+ expected = {'mugwort': 'none', 'birch': 'none', 'alder': 'none', 'ash': 'none', 'oak': 'none',
+ 'grasses': 'purple', 'hazel': 'none'}
assert result == expected
diff --git a/tests/test_sensors.py b/tests/test_sensors.py
index 88863fc..f96d5fb 100644
--- a/tests/test_sensors.py
+++ b/tests/test_sensors.py
@@ -7,7 +7,8 @@
from custom_components.irm_kmi import IrmKmiCoordinator
from custom_components.irm_kmi.binary_sensor import IrmKmiWarning
from custom_components.irm_kmi.const import CONF_LANGUAGE_OVERRIDE
-from custom_components.irm_kmi.sensor import IrmKmiNextWarning, IrmKmiNextSunMove
+from custom_components.irm_kmi.sensor import (IrmKmiNextSunMove,
+ IrmKmiNextWarning)
from tests.conftest import get_api_data