Skip to content

Commit

Permalink
Fix pollen parsing
Browse files Browse the repository at this point in the history
  • Loading branch information
jdejaegh committed Jun 29, 2024
1 parent 3404e16 commit d16a08f
Show file tree
Hide file tree
Showing 8 changed files with 67 additions and 147 deletions.
1 change: 1 addition & 0 deletions custom_components/irm_kmi/const.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
97 changes: 17 additions & 80 deletions custom_components/irm_kmi/pollen.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__(
Expand All @@ -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"""
Expand All @@ -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]:
Expand All @@ -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.
Expand All @@ -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
2 changes: 1 addition & 1 deletion custom_components/irm_kmi/sensor.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
42 changes: 41 additions & 1 deletion tests/fixtures/pollen.svg
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
45 changes: 0 additions & 45 deletions tests/fixtures/pollen_three.svg

This file was deleted.

Loading

0 comments on commit d16a08f

Please sign in to comment.