Improve typing in roomba integration (#102350)

homeassistant/components/roomba/__init__.py

@@ -1,9 +1,11 @@
 """The roomba component."""
 import asyncio
+import contextlib
 from functools import partial
 import logging
+from typing import Any
 
-from roombapy import RoombaConnectionError, RoombaFactory
+from roombapy import Roomba, RoombaConnectionError, RoombaFactory
 
 from homeassistant import exceptions
 from homeassistant.config_entries import ConfigEntry
@@ -16,15 +18,8 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 
-from .const import (
-    BLID,
-    CANCEL_STOP,
-    CONF_BLID,
-    CONF_CONTINUOUS,
-    DOMAIN,
-    PLATFORMS,
-    ROOMBA_SESSION,
-)
+from .const import CONF_BLID, CONF_CONTINUOUS, DOMAIN, PLATFORMS, ROOMBA_SESSION
+from .models import RoombaData
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -62,16 +57,12 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     async def _async_disconnect_roomba(event):
         await async_disconnect_or_timeout(hass, roomba)
 
-    cancel_stop = hass.bus.async_listen_once(
-        EVENT_HOMEASSISTANT_STOP, _async_disconnect_roomba
+    config_entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_disconnect_roomba)
     )
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][config_entry.entry_id] = {
-        ROOMBA_SESSION: roomba,
-        BLID: config_entry.data[CONF_BLID],
-        CANCEL_STOP: cancel_stop,
-    }
+    domain_data = RoombaData(roomba, config_entry.data[CONF_BLID])
+    hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = domain_data
 
     await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
 
@@ -81,7 +72,9 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     return True
 
 
-async def async_connect_or_timeout(hass, roomba):
+async def async_connect_or_timeout(
+    hass: HomeAssistant, roomba: Roomba
+) -> dict[str, Any]:
     """Connect to vacuum."""
     try:
         name = None
@@ -106,12 +99,12 @@ async def async_connect_or_timeout(hass, roomba):
     return {ROOMBA_SESSION: roomba, CONF_NAME: name}
 
 
-async def async_disconnect_or_timeout(hass, roomba):
+async def async_disconnect_or_timeout(hass: HomeAssistant, roomba: Roomba) -> None:
     """Disconnect to vacuum."""
     _LOGGER.debug("Disconnect vacuum")
+    with contextlib.suppress(asyncio.TimeoutError):
         async with asyncio.timeout(3):
             await hass.async_add_executor_job(roomba.disconnect)
-    return True
 
 
 async def async_update_options(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
@@ -125,15 +118,14 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         config_entry, PLATFORMS
     )
     if unload_ok:
-        domain_data = hass.data[DOMAIN][config_entry.entry_id]
-        domain_data[CANCEL_STOP]()
-        await async_disconnect_or_timeout(hass, roomba=domain_data[ROOMBA_SESSION])
+        domain_data: RoombaData = hass.data[DOMAIN][config_entry.entry_id]
+        await async_disconnect_or_timeout(hass, roomba=domain_data.roomba)
         hass.data[DOMAIN].pop(config_entry.entry_id)
 
     return unload_ok
 
 
-def roomba_reported_state(roomba):
+def roomba_reported_state(roomba: Roomba) -> dict[str, Any]:
     """Roomba report."""
     return roomba.master_state.get("state", {}).get("reported", {})
 

homeassistant/components/roomba/binary_sensor.py

@@ -5,8 +5,9 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import roomba_reported_state
-from .const import BLID, DOMAIN, ROOMBA_SESSION
+from .const import DOMAIN
 from .irobot_base import IRobotEntity
+from .models import RoombaData
 
 
 async def async_setup_entry(
@@ -15,9 +16,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the iRobot Roomba vacuum cleaner."""
-    domain_data = hass.data[DOMAIN][config_entry.entry_id]
-    roomba = domain_data[ROOMBA_SESSION]
-    blid = domain_data[BLID]
+    domain_data: RoombaData = hass.data[DOMAIN][config_entry.entry_id]
+    roomba = domain_data.roomba
+    blid = domain_data.blid
     status = roomba_reported_state(roomba).get("bin", {})
     if "full" in status:
         roomba_vac = RoombaBinStatus(roomba, blid)

homeassistant/components/roomba/const.py

@@ -10,5 +10,3 @@ DEFAULT_CERT = "/etc/ssl/certs/ca-certificates.crt"
 DEFAULT_CONTINUOUS = True
 DEFAULT_DELAY = 1
 ROOMBA_SESSION = "roomba_session"
-BLID = "blid_key"
-CANCEL_STOP = "cancel_stop"

homeassistant/components/roomba/models.py

@@ -0,0 +1,14 @@
+"""The roomba integration models."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from roombapy import Roomba
+
+
+@dataclass
+class RoombaData:
+    """Data for the roomba integration."""
+
+    roomba: Roomba
+    blid: str

homeassistant/components/roomba/sensor.py

@@ -16,8 +16,9 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from .const import BLID, DOMAIN, ROOMBA_SESSION
+from .const import DOMAIN
 from .irobot_base import IRobotEntity
+from .models import RoombaData
 
 
 @dataclass
@@ -122,9 +123,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the iRobot Roomba vacuum cleaner."""
-    domain_data = hass.data[DOMAIN][config_entry.entry_id]
-    roomba = domain_data[ROOMBA_SESSION]
-    blid = domain_data[BLID]
+    domain_data: RoombaData = hass.data[DOMAIN][config_entry.entry_id]
+    roomba = domain_data.roomba
+    blid = domain_data.blid
 
     async_add_entities(
         RoombaSensor(roomba, blid, entity_description) for entity_description in SENSORS

homeassistant/components/roomba/vacuum.py

@@ -7,8 +7,9 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import roomba_reported_state
 from .braava import BraavaJet
-from .const import BLID, DOMAIN, ROOMBA_SESSION
+from .const import DOMAIN
 from .irobot_base import IRobotVacuum
+from .models import RoombaData
 from .roomba import RoombaVacuum, RoombaVacuumCarpetBoost
 
 
@@ -18,9 +19,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the iRobot Roomba vacuum cleaner."""
-    domain_data = hass.data[DOMAIN][config_entry.entry_id]
-    roomba = domain_data[ROOMBA_SESSION]
-    blid = domain_data[BLID]
+    domain_data: RoombaData = hass.data[DOMAIN][config_entry.entry_id]
+    roomba = domain_data.roomba
+    blid = domain_data.blid
 
     # Get the capabilities of our unit
     state = roomba_reported_state(roomba)