ha-noaa-space-weather/custom-component/space_weather/sensor.py

import logging
from datetime import timedelta, datetime

import aiohttp
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.entity import Entity
from homeassistant.util import Throttle

_LOGGER = logging.getLogger(__name__)

SCALES_URL = 'https://services.swpc.noaa.gov/products/noaa-scales.json'
SCAN_INTERVAL = timedelta(minutes=5)


async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
    session = async_get_clientsession(hass)
    async_add_entities([
        SpaceWeatherScaleSensor(session, 'R', '0', None),
        SpaceWeatherScaleSensor(session, 'S', '0', None),
        SpaceWeatherScaleSensor(session, 'G', '0', None),
        SpaceWeatherScaleSensor(session, 'R', '-1', '24hr_max'),
        SpaceWeatherScaleSensor(session, 'S', '-1', '24hr_max'),
        SpaceWeatherScaleSensor(session, 'G', '-1', '24hr_max'),

        SpaceWeatherPredictionSensor(session, 'R', 'MinorProb', '1', 'today'),
        SpaceWeatherPredictionSensor(session, 'R', 'MajorProb', '1', 'today'),
        SpaceWeatherPredictionSensor(session, 'S', 'Scale', '1', 'today'),
        SpaceWeatherPredictionSensor(session, 'S', 'Prob', '1', 'today'),
        SpaceWeatherPredictionSensor(session, 'G', 'Scale', '1', 'today'),
        SpaceWeatherPredictionSensor(session, 'R', 'MinorProb', '2', '1day'),
        SpaceWeatherPredictionSensor(session, 'R', 'MajorProb', '2', '1day'),
        SpaceWeatherPredictionSensor(session, 'S', 'Scale', '2', '1day'),
        SpaceWeatherPredictionSensor(session, 'S', 'Prob', '2', '1day'),
        SpaceWeatherPredictionSensor(session, 'G', 'Scale', '2', '1day'),
        SpaceWeatherPredictionSensor(session, 'R', 'MinorProb', '3', '2day'),
        SpaceWeatherPredictionSensor(session, 'R', 'MajorProb', '3', '2day'),
        SpaceWeatherPredictionSensor(session, 'S', 'Scale', '3', '2day'),
        SpaceWeatherPredictionSensor(session, 'S', 'Prob', '3', '2day'),
        SpaceWeatherPredictionSensor(session, 'G', 'Scale', '3', '2day'),

        PlanetaryKIndexSensor(session)
    ], True)


class SpaceWeatherScaleSensor(Entity):
    def __init__(self, session, scale_key, data_selector, trailing):
        self._session = session
        self._scale_key = scale_key

        self._name = f'Space Weather Scale {scale_key}'
        if trailing is not None and len(trailing):
            self._name = self._name + ' ' + trailing.replace('_', ' ').replace('  ', ' ')

        self._data = None
        assert isinstance(data_selector, str)
        self._data_selector = data_selector
        self._trailing = trailing

    @property
    def name(self):
        return self._name

    @property
    def unique_id(self):
        s = f'space_weather_scale_{self._scale_key.lower()}'
        if self._trailing is not None and len(self._trailing):
            s = s + '_' + self._trailing.strip('_')
        return s

    @property
    def state(self):
        return f'{self._scale_key}{self._data[self._scale_key]["Scale"]}'

    @property
    def extra_state_attributes(self):
        if self._data:
            return {
                'scale_int': int(self._data[self._scale_key]['Scale']),
                'text': self._data[self._scale_key]['Text'],
                'timestamp': datetime.fromisoformat(self._data["DateStamp"] + 'T' + self._data["TimeStamp"] + '+00:00').isoformat()
            }
        return None

    @Throttle(SCAN_INTERVAL)
    async def async_update(self):
        try:
            async with self._session.get(SCALES_URL) as response:
                if response.status == 200:
                    data = await response.json()
                    self._data = data[self._data_selector]
                else:
                    _LOGGER.error(f'Error fetching data from {SCALES_URL}')
        except aiohttp.ClientError as err:
            _LOGGER.error(f'Error fetching data from {SCALES_URL}: {err}')


class SpaceWeatherPredictionSensor(Entity):
    def __init__(self, session, scale_key, pred_key, data_selector, trailing):
        self._session = session
        self._scale_key = scale_key
        self._pred_key = pred_key
        self._data_selector = data_selector
        self._trailing = trailing
        self._name = f'Space Weather Prediction {scale_key} {pred_key} {trailing.replace("_", " ").replace("  ", " ")}'
        self._state = None
        self._data = None

    @property
    def name(self):
        return self._name

    @property
    def unique_id(self):
        return f'space_weather_pred_{self._scale_key}_{self._pred_key}_{self._trailing}'.lower()

    @property
    def state(self):
        if self._pred_key == 'Scale':
            return self._state
        elif self._state is not None:
            try:
                return float(self._state)
            except ValueError:
                return None
        return None

    @property
    def unit_of_measurement(self):
        if self._pred_key in ['MinorProb', 'MajorProb', 'Prob']:
            return '%'
        return None

    @property
    def extra_state_attributes(self):
        if self._data:
            return {
                'timestamp': datetime.fromisoformat(self._data["DateStamp"] + 'T' + self._data["TimeStamp"] + '+00:00').isoformat()
            }
        return None

    @Throttle(SCAN_INTERVAL)
    async def async_update(self):
        try:
            async with self._session.get(SCALES_URL) as response:
                if response.status == 200:
                    data = await response.json()
                    self._data = data[self._data_selector]
                    self._state = self._data[self._scale_key][self._pred_key]
                else:
                    _LOGGER.error(f'Error fetching data from {SCALES_URL}')
        except aiohttp.ClientError as err:
            _LOGGER.error(f'Error fetching data from {SCALES_URL}: {err}')


class PlanetaryKIndexSensor(Entity):
    """
    Temporary sensor until https://github.com/tcarwash/home-assistant_noaa-space-weather gets their shit together
    """

    def __init__(self, session):
        self._session = session
        self._name = 'Space Weather Planetary K-Index'
        self._data = None

    @property
    def name(self):
        return self._name

    @property
    def unique_id(self):
        return 'space_weather_planetary_k_index'

    @property
    def state(self):
        if self._data:
            return float(self._data['kp_index'])
        return None

    # @property
    # def device_class(self):
    #     return None

    @property
    def extra_state_attributes(self):
        if self._data:
            return {
                'kp_index': float(self._data['kp_index']),
                'estimated_kp': float(self._data['estimated_kp']),
                'timestamp': datetime.fromisoformat(self._data['time_tag']).isoformat(),
                'state_class': 'measurement'
            }
        return None

    @Throttle(SCAN_INTERVAL)
    async def async_update(self):
        try:
            async with self._session.get('https://services.swpc.noaa.gov/json/planetary_k_index_1m.json') as response:
                if response.status == 200:
                    data = await response.json()
                    self._data = data[-1]
                else:
                    _LOGGER.error(f'Error fetching data from {SCALES_URL}')
        except aiohttp.ClientError as err:
            _LOGGER.error(f'Error fetching data from {SCALES_URL}: {err}')