Add full typing to kraken (#50718)

* Add full typing to kraken * Let device_info return DeviceInfo * Replace unsub_listeners with entry.async_on_unload * Raise TypeError on end of _try_get_state * Assert Coordinator is not none * Add class SensorType * Add strict typing to kraken * Add changes from code review * Revert typed dict creation
2021-05-17 09:12:04 +02:00 · 2021-05-17 09:12:04 +02:00 · 663c0374ab
commit 663c0374ab
parent 120bf8aed7
7 changed files with 137 additions and 80 deletions
--- a/.strict-typing
+++ b/.strict-typing
@ -29,6 +29,7 @@ homeassistant.components.hyperion.*
 homeassistant.components.image_processing.*
 homeassistant.components.integration.*
 homeassistant.components.knx.*
+homeassistant.components.kraken.*
 homeassistant.components.light.*
 homeassistant.components.lock.*
 homeassistant.components.mailbox.*
--- a/homeassistant/components/kraken/init.py
+++ b/homeassistant/components/kraken/init.py
@ -21,6 +21,7 @@ from .const import (
    DEFAULT_TRACKED_ASSET_PAIR,
    DISPATCH_CONFIG_UPDATED,
    DOMAIN,
+    KrakenResponse,
 )
 from .utils import get_tradable_asset_pairs

@ -47,8 +48,6 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
        config_entry, PLATFORMS
    )
    if unload_ok:
-        for unsub_listener in hass.data[DOMAIN].unsub_listeners:
-            unsub_listener()
        hass.data.pop(DOMAIN)

    return unload_ok
@ -62,11 +61,10 @@ class KrakenData:
        self._hass = hass
        self._config_entry = config_entry
        self._api = pykrakenapi.KrakenAPI(krakenex.API(), retry=0, crl_sleep=0)
-        self.tradable_asset_pairs = None
-        self.coordinator = None
-        self.unsub_listeners = []
+        self.tradable_asset_pairs: dict[str, str] = {}
+        self.coordinator: DataUpdateCoordinator[KrakenResponse | None] | None = None

-    async def async_update(self) -> None:
+    async def async_update(self) -> KrakenResponse | None:
        """Get the latest data from the Kraken.com REST API.

        All tradeable asset pairs are retrieved, not the tracked asset pairs
@ -91,8 +89,9 @@ class KrakenData:
            _LOGGER.warning(
                "Exceeded the Kraken.com call rate limit. Increase the update interval to prevent this error"
            )
+        return None

-    def _get_kraken_data(self) -> dict:
+    def _get_kraken_data(self) -> KrakenResponse:
        websocket_name_pairs = self._get_websocket_name_asset_pairs()
        ticker_df = self._api.get_ticker_information(websocket_name_pairs)
        # Rename columns to their full name
@ -109,7 +108,7 @@ class KrakenData:
                "o": "opening_price",
            }
        )
-        response_dict = ticker_df.transpose().to_dict()
+        response_dict: KrakenResponse = ticker_df.transpose().to_dict()
        return response_dict

    async def _async_refresh_tradable_asset_pairs(self) -> None:
@ -140,12 +139,13 @@ class KrakenData:
        )
        await self.coordinator.async_config_entry_first_refresh()

-    def _get_websocket_name_asset_pairs(self) -> list:
+    def _get_websocket_name_asset_pairs(self) -> str:
        return ",".join(wsname for wsname in self.tradable_asset_pairs.values())

    def set_update_interval(self, update_interval: int) -> None:
        """Set the coordinator update_interval to the supplied update_interval."""
-        self.coordinator.update_interval = timedelta(seconds=update_interval)
+        if self.coordinator is not None:
+            self.coordinator.update_interval = timedelta(seconds=update_interval)


 async def async_options_updated(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
--- a/homeassistant/components/kraken/config_flow.py
+++ b/homeassistant/components/kraken/config_flow.py
@ -1,5 +1,8 @@
 """Config flow for kraken integration."""
+from __future__ import annotations
+
 import logging
+from typing import Any

 import krakenex
 from pykrakenapi.pykrakenapi import KrakenAPI
@ -8,6 +11,7 @@ import voluptuous as vol
 from homeassistant import config_entries
 from homeassistant.const import CONF_SCAN_INTERVAL
 from homeassistant.core import callback
+from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers import config_validation as cv

 from .const import CONF_TRACKED_ASSET_PAIRS, DEFAULT_SCAN_INTERVAL, DOMAIN
@ -24,11 +28,15 @@ class KrakenConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):

    @staticmethod
    @callback
-    def async_get_options_flow(config_entry):
+    def async_get_options_flow(
+        config_entry: config_entries.ConfigEntry,
+    ) -> KrakenOptionsFlowHandler:
        """Get the options flow for this handler."""
        return KrakenOptionsFlowHandler(config_entry)

-    async def async_step_user(self, user_input=None):
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
        """Handle the initial step."""
        if DOMAIN in self.hass.data:
            return self.async_abort(reason="already_configured")
@ -44,11 +52,13 @@ class KrakenConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 class KrakenOptionsFlowHandler(config_entries.OptionsFlow):
    """Handle Kraken client options."""

-    def __init__(self, config_entry):
+    def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
        """Initialize Kraken options flow."""
        self.config_entry = config_entry

-    async def async_step_init(self, user_input=None):
+    async def async_step_init(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
        """Manage the Kraken options."""
        if user_input is not None:
            return self.async_create_entry(title="", data=user_input)
--- a/homeassistant/components/kraken/const.py
+++ b/homeassistant/components/kraken/const.py
@ -1,5 +1,19 @@
 """Constants for the kraken integration."""

+from __future__ import annotations
+
+from typing import Dict, TypedDict
+
+KrakenResponse = Dict[str, Dict[str, float]]
+
+
+class SensorType(TypedDict):
+    """SensorType class."""
+
+    name: str
+    enabled_by_default: bool
+
+
 DEFAULT_SCAN_INTERVAL = 60
 DEFAULT_TRACKED_ASSET_PAIR = "XBT/USD"
 DISPATCH_CONFIG_UPDATED = "kraken_config_updated"
@ -8,7 +22,7 @@ CONF_TRACKED_ASSET_PAIRS = "tracked_asset_pairs"

 DOMAIN = "kraken"

-SENSOR_TYPES = [
+SENSOR_TYPES: list[SensorType] = [
    {"name": "ask", "enabled_by_default": True},
    {"name": "ask_volume", "enabled_by_default": False},
    {"name": "bid", "enabled_by_default": True},
--- a/homeassistant/components/kraken/sensor.py
+++ b/homeassistant/components/kraken/sensor.py
@ -8,6 +8,9 @@ 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
@ -16,12 +19,17 @@ from .const import (
    DISPATCH_CONFIG_UPDATED,
    DOMAIN,
    SENSOR_TYPES,
+    SensorType,
 )

 _LOGGER = logging.getLogger(__name__)


-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
    """Add kraken entities from a config_entry."""

    @callback
@ -59,7 +67,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):

    async_update_sensors(hass, config_entry)

-    hass.data[DOMAIN].unsub_listeners.append(
+    config_entry.async_on_unload(
        async_dispatcher_connect(
            hass,
            DISPATCH_CONFIG_UPDATED,
@ -75,9 +83,10 @@ class KrakenSensor(CoordinatorEntity, SensorEntity):
        self,
        kraken_data: KrakenData,
        tracked_asset_pair: str,
-        sensor_type: dict[str, bool],
+        sensor_type: SensorType,
    ) -> None:
        """Initialize."""
+        assert kraken_data.coordinator is not None
        super().__init__(kraken_data.coordinator)
        self.tracked_asset_pair_wsname = kraken_data.tradable_asset_pairs[
            tracked_asset_pair
@ -100,22 +109,22 @@ class KrakenSensor(CoordinatorEntity, SensorEntity):
        self._state = None

    @property
-    def entity_registry_enabled_default(self):
+    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):
+    def name(self) -> str:
        """Return the name."""
        return self._name

    @property
-    def unique_id(self):
+    def unique_id(self) -> str:
        """Set unique_id for sensor."""
        return self._name.lower()

    @property
-    def state(self):
+    def state(self) -> StateType:
        """Return the state."""
        return self._state

@ -124,13 +133,76 @@ class KrakenSensor(CoordinatorEntity, SensorEntity):
        await super().async_added_to_hass()
        self._update_internal_state()

-    def _handle_coordinator_update(self):
+    def _handle_coordinator_update(self) -> None:
        self._update_internal_state()
        super()._handle_coordinator_update()

-    def _update_internal_state(self):
+    def _update_internal_state(self) -> None:
        try:
-            self._state = self._try_get_state()
+            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._received_data_at_least_once = True  # Received data at least one time.
        except TypeError:
            if self._received_data_at_least_once:
@ -141,55 +213,8 @@ class KrakenSensor(CoordinatorEntity, SensorEntity):
                    )
                    self._available = False

-    def _try_get_state(self) -> str:
-        """Try to get the state or return a TypeError."""
-        if self._sensor_type == "last_trade_closed":
-            return self.coordinator.data[self.tracked_asset_pair_wsname][
-                "last_trade_closed"
-            ][0]
-        if self._sensor_type == "ask":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["ask"][0]
-        if self._sensor_type == "ask_volume":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["ask"][1]
-        if self._sensor_type == "bid":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["bid"][0]
-        if self._sensor_type == "bid_volume":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["bid"][1]
-        if self._sensor_type == "volume_today":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["volume"][0]
-        if self._sensor_type == "volume_last_24h":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["volume"][1]
-        if self._sensor_type == "volume_weighted_average_today":
-            return self.coordinator.data[self.tracked_asset_pair_wsname][
-                "volume_weighted_average"
-            ][0]
-        if self._sensor_type == "volume_weighted_average_last_24h":
-            return self.coordinator.data[self.tracked_asset_pair_wsname][
-                "volume_weighted_average"
-            ][1]
-        if self._sensor_type == "number_of_trades_today":
-            return self.coordinator.data[self.tracked_asset_pair_wsname][
-                "number_of_trades"
-            ][0]
-        if self._sensor_type == "number_of_trades_last_24h":
-            return self.coordinator.data[self.tracked_asset_pair_wsname][
-                "number_of_trades"
-            ][1]
-        if self._sensor_type == "low_today":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["low"][0]
-        if self._sensor_type == "low_last_24h":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["low"][1]
-        if self._sensor_type == "high_today":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["high"][0]
-        if self._sensor_type == "high_last_24h":
-            return self.coordinator.data[self.tracked_asset_pair_wsname]["high"][1]
-        if self._sensor_type == "opening_price_today":
-            return self.coordinator.data[self.tracked_asset_pair_wsname][
-                "opening_price"
-            ]
-
    @property
-    def icon(self):
+    def icon(self) -> str:
        """Return the icon."""
        if self._target_asset == "EUR":
            return "mdi:currency-eur"
@ -204,19 +229,19 @@ class KrakenSensor(CoordinatorEntity, SensorEntity):
        return "mdi:cash"

    @property
-    def unit_of_measurement(self):
+    def 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):
+    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) -> dict:
+    def device_info(self) -> DeviceInfo:
        """Return a device description for device registry."""

        return {
--- a/mypy.ini
+++ b/mypy.ini
@ -330,6 +330,17 @@ no_implicit_optional = true
 warn_return_any = true
 warn_unreachable = true

+[mypy-homeassistant.components.kraken.*]
+check_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+warn_return_any = true
+warn_unreachable = true
+
 [mypy-homeassistant.components.light.*]
 check_untyped_defs = true
 disallow_incomplete_defs = true
@ -982,9 +993,6 @@ ignore_errors = true
 [mypy-homeassistant.components.kostal_plenticore.*]
 ignore_errors = true

-[mypy-homeassistant.components.kraken.*]
-ignore_errors = true
-
 [mypy-homeassistant.components.kulersky.*]
 ignore_errors = true

--- a/script/hassfest/mypy_config.py
+++ b/script/hassfest/mypy_config.py
@ -114,7 +114,6 @@ IGNORED_MODULES: Final[list[str]] = [
    "homeassistant.components.kodi.*",
    "homeassistant.components.konnected.*",
    "homeassistant.components.kostal_plenticore.*",
-    "homeassistant.components.kraken.*",
    "homeassistant.components.kulersky.*",
    "homeassistant.components.lifx.*",
    "homeassistant.components.litejet.*",