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

.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.*

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.tradable_asset_pairs: dict[str, str] = {}
-        self.coordinator = None
+        self.coordinator: DataUpdateCoordinator[KrakenResponse | None] | None = None
-        self.unsub_listeners = []
 
-    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:

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)

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},

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 {

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
 

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.*",