From b634bd26d034169825dea58b67d2d2410c835791 Mon Sep 17 00:00:00 2001 From: Marc Mueller <30130371+cdce8p@users.noreply.github.com> Date: Thu, 23 Sep 2021 16:42:55 +0200 Subject: [PATCH] Use EntityDescription - kraken (#56436) --- homeassistant/components/kraken/const.py | 134 +++++++++++++++---- homeassistant/components/kraken/sensor.py | 152 +++++----------------- 2 files changed, 141 insertions(+), 145 deletions(-) diff --git a/homeassistant/components/kraken/const.py b/homeassistant/components/kraken/const.py index 2272d12ead6..669d64a49c8 100644 --- a/homeassistant/components/kraken/const.py +++ b/homeassistant/components/kraken/const.py @@ -1,17 +1,28 @@ """Constants for the kraken integration.""" - from __future__ import annotations -from typing import Dict, TypedDict +from dataclasses import dataclass +from typing import Callable, Dict, TypedDict -KrakenResponse = Dict[str, Dict[str, float]] +from homeassistant.components.sensor import SensorEntityDescription +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -class SensorType(TypedDict): - """SensorType class.""" +class KrakenResponseEntry(TypedDict): + """Dict describing a single response entry.""" - name: str - enabled_by_default: bool + ask: tuple[float, float, float] + bid: tuple[float, float, float] + last_trade_closed: tuple[float, float] + volume: tuple[float, float] + volume_weighted_average: tuple[float, float] + number_of_trades: tuple[int, int] + low: tuple[float, float] + high: tuple[float, float] + opening_price: float + + +KrakenResponse = Dict[str, KrakenResponseEntry] DEFAULT_SCAN_INTERVAL = 60 @@ -22,21 +33,94 @@ CONF_TRACKED_ASSET_PAIRS = "tracked_asset_pairs" DOMAIN = "kraken" -SENSOR_TYPES: list[SensorType] = [ - {"name": "ask", "enabled_by_default": True}, - {"name": "ask_volume", "enabled_by_default": False}, - {"name": "bid", "enabled_by_default": True}, - {"name": "bid_volume", "enabled_by_default": False}, - {"name": "volume_today", "enabled_by_default": False}, - {"name": "volume_last_24h", "enabled_by_default": False}, - {"name": "volume_weighted_average_today", "enabled_by_default": False}, - {"name": "volume_weighted_average_last_24h", "enabled_by_default": False}, - {"name": "number_of_trades_today", "enabled_by_default": False}, - {"name": "number_of_trades_last_24h", "enabled_by_default": False}, - {"name": "last_trade_closed", "enabled_by_default": False}, - {"name": "low_today", "enabled_by_default": True}, - {"name": "low_last_24h", "enabled_by_default": False}, - {"name": "high_today", "enabled_by_default": True}, - {"name": "high_last_24h", "enabled_by_default": False}, - {"name": "opening_price_today", "enabled_by_default": False}, -] + +@dataclass +class KrakenRequiredKeysMixin: + """Mixin for required keys.""" + + value_fn: Callable[[DataUpdateCoordinator[KrakenResponse], str], float | int] + + +@dataclass +class KrakenSensorEntityDescription(SensorEntityDescription, KrakenRequiredKeysMixin): + """Describes Kraken sensor entity.""" + + +SENSOR_TYPES: tuple[KrakenSensorEntityDescription, ...] = ( + KrakenSensorEntityDescription( + key="ask", + value_fn=lambda x, y: x.data[y]["ask"][0], + ), + KrakenSensorEntityDescription( + key="ask_volume", + value_fn=lambda x, y: x.data[y]["ask"][1], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="bid", + value_fn=lambda x, y: x.data[y]["bid"][0], + ), + KrakenSensorEntityDescription( + key="bid_volume", + value_fn=lambda x, y: x.data[y]["bid"][1], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="volume_today", + value_fn=lambda x, y: x.data[y]["volume"][0], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="volume_last_24h", + value_fn=lambda x, y: x.data[y]["volume"][1], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="volume_weighted_average_today", + value_fn=lambda x, y: x.data[y]["volume_weighted_average"][0], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="volume_weighted_average_last_24h", + value_fn=lambda x, y: x.data[y]["volume_weighted_average"][1], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="number_of_trades_today", + value_fn=lambda x, y: x.data[y]["number_of_trades"][0], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="number_of_trades_last_24h", + value_fn=lambda x, y: x.data[y]["number_of_trades"][1], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="last_trade_closed", + value_fn=lambda x, y: x.data[y]["last_trade_closed"][0], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="low_today", + value_fn=lambda x, y: x.data[y]["low"][0], + ), + KrakenSensorEntityDescription( + key="low_last_24h", + value_fn=lambda x, y: x.data[y]["low"][1], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="high_today", + value_fn=lambda x, y: x.data[y]["high"][0], + ), + KrakenSensorEntityDescription( + key="high_last_24h", + value_fn=lambda x, y: x.data[y]["high"][1], + entity_registry_enabled_default=False, + ), + KrakenSensorEntityDescription( + key="opening_price_today", + value_fn=lambda x, y: x.data[y]["opening_price"], + entity_registry_enabled_default=False, + ), +) diff --git a/homeassistant/components/kraken/sensor.py b/homeassistant/components/kraken/sensor.py index 1b9f8ca13cc..9c2030766f7 100644 --- a/homeassistant/components/kraken/sensor.py +++ b/homeassistant/components/kraken/sensor.py @@ -2,15 +2,14 @@ from __future__ import annotations import logging +from typing import Optional from homeassistant.components.sensor import SensorEntity from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import device_registry from homeassistant.helpers.dispatcher import async_dispatcher_connect -from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity from . import KrakenData @@ -19,7 +18,8 @@ from .const import ( DISPATCH_CONFIG_UPDATED, DOMAIN, SENSOR_TYPES, - SensorType, + KrakenResponse, + KrakenSensorEntityDescription, ) _LOGGER = logging.getLogger(__name__) @@ -43,7 +43,7 @@ async def async_setup_entry( ) } - sensors = [] + entities = [] for tracked_asset_pair in config_entry.options[CONF_TRACKED_ASSET_PAIRS]: # Only create new devices if ( @@ -51,15 +51,17 @@ async def async_setup_entry( ) in existing_devices: existing_devices.pop(device_name) else: - for sensor_type in SENSOR_TYPES: - sensors.append( + entities.extend( + [ KrakenSensor( hass.data[DOMAIN], tracked_asset_pair, - sensor_type, + description, ) - ) - async_add_entities(sensors, True) + for description in SENSOR_TYPES + ] + ) + async_add_entities(entities, True) # Remove devices for asset pairs which are no longer tracked for device_id in existing_devices.values(): @@ -76,57 +78,46 @@ async def async_setup_entry( ) -class KrakenSensor(CoordinatorEntity, SensorEntity): +class KrakenSensor(CoordinatorEntity[Optional[KrakenResponse]], SensorEntity): """Define a Kraken sensor.""" + entity_description: KrakenSensorEntityDescription + def __init__( self, kraken_data: KrakenData, tracked_asset_pair: str, - sensor_type: SensorType, + description: KrakenSensorEntityDescription, ) -> None: """Initialize.""" assert kraken_data.coordinator is not None super().__init__(kraken_data.coordinator) + self.entity_description = description self.tracked_asset_pair_wsname = kraken_data.tradable_asset_pairs[ tracked_asset_pair ] - self._source_asset = tracked_asset_pair.split("/")[0] + source_asset = tracked_asset_pair.split("/")[0] self._target_asset = tracked_asset_pair.split("/")[1] - self._sensor_type = sensor_type["name"] - self._enabled_by_default = sensor_type["enabled_by_default"] - self._unit_of_measurement = self._target_asset - self._device_name = f"{self._source_asset} {self._target_asset}" - self._name = "_".join( + if "number_of" not in description.key: + self._attr_native_unit_of_measurement = self._target_asset + self._device_name = f"{source_asset} {self._target_asset}" + self._attr_name = "_".join( [ tracked_asset_pair.split("/")[0], tracked_asset_pair.split("/")[1], - sensor_type["name"], + description.key, ] ) + self._attr_unique_id = self._attr_name.lower() self._received_data_at_least_once = False self._available = True - self._state = None - @property - def entity_registry_enabled_default(self) -> bool: - """Return if the entity should be enabled when first added to the entity registry.""" - return self._enabled_by_default - - @property - def name(self) -> str: - """Return the name.""" - return self._name - - @property - def unique_id(self) -> str: - """Set unique_id for sensor.""" - return self._name.lower() - - @property - def native_value(self) -> StateType: - """Return the state.""" - return self._state + self._attr_device_info = { + "identifiers": {(DOMAIN, f"{source_asset}_{self._target_asset}")}, + "name": self._device_name, + "manufacturer": "Kraken.com", + "entry_type": "service", + } async def async_added_to_hass(self) -> None: """Handle entity which will be added.""" @@ -139,70 +130,9 @@ class KrakenSensor(CoordinatorEntity, SensorEntity): def _update_internal_state(self) -> None: try: - if self._sensor_type == "last_trade_closed": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "last_trade_closed" - ][0] - if self._sensor_type == "ask": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "ask" - ][0] - if self._sensor_type == "ask_volume": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "ask" - ][1] - if self._sensor_type == "bid": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "bid" - ][0] - if self._sensor_type == "bid_volume": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "bid" - ][1] - if self._sensor_type == "volume_today": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "volume" - ][0] - if self._sensor_type == "volume_last_24h": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "volume" - ][1] - if self._sensor_type == "volume_weighted_average_today": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "volume_weighted_average" - ][0] - if self._sensor_type == "volume_weighted_average_last_24h": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "volume_weighted_average" - ][1] - if self._sensor_type == "number_of_trades_today": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "number_of_trades" - ][0] - if self._sensor_type == "number_of_trades_last_24h": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "number_of_trades" - ][1] - if self._sensor_type == "low_today": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "low" - ][0] - if self._sensor_type == "low_last_24h": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "low" - ][1] - if self._sensor_type == "high_today": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "high" - ][0] - if self._sensor_type == "high_last_24h": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "high" - ][1] - if self._sensor_type == "opening_price_today": - self._state = self.coordinator.data[self.tracked_asset_pair_wsname][ - "opening_price" - ] + self._attr_native_value = self.entity_description.value_fn( + self.coordinator, self.tracked_asset_pair_wsname # type: ignore[arg-type] + ) self._received_data_at_least_once = True # Received data at least one time. except TypeError: if self._received_data_at_least_once: @@ -228,29 +158,11 @@ class KrakenSensor(CoordinatorEntity, SensorEntity): return "mdi:currency-btc" return "mdi:cash" - @property - def native_unit_of_measurement(self) -> str | None: - """Return the unit the value is expressed in.""" - if "number_of" not in self._sensor_type: - return self._unit_of_measurement - return None - @property def available(self) -> bool: """Could the api be accessed during the last update call.""" return self._available and self.coordinator.last_update_success - @property - def device_info(self) -> DeviceInfo: - """Return a device description for device registry.""" - - return { - "identifiers": {(DOMAIN, f"{self._source_asset}_{self._target_asset}")}, - "name": self._device_name, - "manufacturer": "Kraken.com", - "entry_type": "service", - } - def create_device_name(tracked_asset_pair: str) -> str: """Create the device name for a given tracked asset pair."""