From 7cb01f75ae709961ede7bb8c91826637f937b441 Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Tue, 2 Apr 2024 20:21:55 +0200 Subject: [PATCH] Add typing to Roomba config flow (#114624) --- .../components/roomba/config_flow.py | 106 +++++++++--------- 1 file changed, 55 insertions(+), 51 deletions(-) diff --git a/homeassistant/components/roomba/config_flow.py b/homeassistant/components/roomba/config_flow.py index 7b834421135..53ea9aa7c44 100644 --- a/homeassistant/components/roomba/config_flow.py +++ b/homeassistant/components/roomba/config_flow.py @@ -4,8 +4,9 @@ from __future__ import annotations import asyncio from functools import partial +from typing import Any -from roombapy import RoombaFactory +from roombapy import RoombaFactory, RoombaInfo from roombapy.discovery import RoombaDiscovery from roombapy.getpassword import RoombaPassword import voluptuous as vol @@ -15,7 +16,7 @@ from homeassistant.config_entries import ( ConfigEntry, ConfigFlow, ConfigFlowResult, - OptionsFlow, + OptionsFlowWithConfigEntry, ) from homeassistant.const import CONF_DELAY, CONF_HOST, CONF_NAME, CONF_PASSWORD from homeassistant.core import HomeAssistant, callback @@ -43,7 +44,7 @@ AUTH_HELP_URL_KEY = "auth_help_url" AUTH_HELP_URL_VALUE = "https://www.home-assistant.io/integrations/roomba/#manually-retrieving-your-credentials" -async def validate_input(hass: HomeAssistant, data): +async def validate_input(hass: HomeAssistant, data: dict[str, Any]) -> dict[str, Any]: """Validate the user input allows us to connect. Data has the keys from DATA_SCHEMA with values provided by the user. @@ -75,20 +76,21 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): VERSION = 1 - def __init__(self): + name: str | None = None + blid: str | None = None + host: str | None = None + + def __init__(self) -> None: """Initialize the roomba flow.""" - self.discovered_robots = {} - self.name = None - self.blid = None - self.host = None + self.discovered_robots: dict[str, RoombaInfo] = {} @staticmethod @callback def async_get_options_flow( config_entry: ConfigEntry, - ) -> OptionsFlowHandler: + ) -> RoombaOptionsFlowHandler: """Get the options flow for this handler.""" - return OptionsFlowHandler(config_entry) + return RoombaOptionsFlowHandler(config_entry) async def async_step_zeroconf( self, discovery_info: zeroconf.ZeroconfServiceInfo @@ -135,8 +137,9 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): self.context["title_placeholders"] = {"host": self.host, "name": self.blid} return await self.async_step_user() - async def _async_start_link(self): + async def _async_start_link(self) -> ConfigFlowResult: """Start linking.""" + assert self.host device = self.discovered_robots[self.host] self.blid = device.blid self.name = device.robot_name @@ -144,7 +147,9 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): self._abort_if_unique_id_configured() return await self.async_step_link() - async def async_step_user(self, user_input=None): + async def async_step_user( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: """Handle a flow start.""" # Check if user chooses manual entry if user_input is not None and not user_input.get(CONF_HOST): @@ -181,25 +186,23 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): if not self.discovered_robots: return await self.async_step_manual() + hosts: dict[str | None, str] = { + **{ + device.ip: f"{device.robot_name} ({device.ip})" + for device in devices + if device.blid not in already_configured + }, + None: "Manually add a Roomba or Braava", + } + return self.async_show_form( step_id="user", - data_schema=vol.Schema( - { - vol.Optional("host"): vol.In( - { - **{ - device.ip: f"{device.robot_name} ({device.ip})" - for device in devices - if device.blid not in already_configured - }, - None: "Manually add a Roomba or Braava", - } - ) - } - ), + data_schema=vol.Schema({vol.Optional("host"): vol.In(hosts)}), ) - async def async_step_manual(self, user_input=None): + async def async_step_manual( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: """Handle manual device setup.""" if user_input is None: return self.async_show_form( @@ -224,7 +227,9 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): self._abort_if_unique_id_configured() return await self.async_step_link() - async def async_step_link(self, user_input=None): + async def async_step_link( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: """Attempt to link with the Roomba. Given a configured host, will ask the user to press the home and target buttons @@ -235,7 +240,7 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): step_id="link", description_placeholders={CONF_NAME: self.name or self.blid}, ) - + assert self.host roomba_pw = RoombaPassword(self.host) try: @@ -260,10 +265,12 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): return self.async_abort(reason="cannot_connect") self.name = info[CONF_NAME] - + assert self.name return self.async_create_entry(title=self.name, data=config) - async def async_step_link_manual(self, user_input=None): + async def async_step_link_manual( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: """Handle manual linking.""" errors = {} @@ -278,8 +285,7 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): info = await validate_input(self.hass, config) except CannotConnect: errors = {"base": "cannot_connect"} - - if not errors: + else: return self.async_create_entry(title=info[CONF_NAME], data=config) return self.async_show_form( @@ -290,14 +296,12 @@ class RoombaConfigFlow(ConfigFlow, domain=DOMAIN): ) -class OptionsFlowHandler(OptionsFlow): +class RoombaOptionsFlowHandler(OptionsFlowWithConfigEntry): """Handle options.""" - def __init__(self, config_entry: ConfigEntry) -> None: - """Initialize 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 + ) -> ConfigFlowResult: """Manage the options.""" if user_input is not None: return self.async_create_entry(title="", data=user_input) @@ -308,15 +312,11 @@ class OptionsFlowHandler(OptionsFlow): { vol.Optional( CONF_CONTINUOUS, - default=self.config_entry.options.get( - CONF_CONTINUOUS, DEFAULT_CONTINUOUS - ), + default=self.options.get(CONF_CONTINUOUS, DEFAULT_CONTINUOUS), ): bool, vol.Optional( CONF_DELAY, - default=self.config_entry.options.get( - CONF_DELAY, DEFAULT_DELAY - ), + default=self.options.get(CONF_DELAY, DEFAULT_DELAY), ): int, } ), @@ -324,7 +324,7 @@ class OptionsFlowHandler(OptionsFlow): @callback -def _async_get_roomba_discovery(): +def _async_get_roomba_discovery() -> RoombaDiscovery: """Create a discovery object.""" discovery = RoombaDiscovery() discovery.amount_of_broadcasted_messages = MAX_NUM_DEVICES_TO_DISCOVER @@ -332,24 +332,28 @@ def _async_get_roomba_discovery(): @callback -def _async_blid_from_hostname(hostname): +def _async_blid_from_hostname(hostname: str) -> str: """Extract the blid from the hostname.""" return hostname.split("-")[1].split(".")[0].upper() -async def _async_discover_roombas(hass, host): - discovered_hosts = set() - devices = [] +async def _async_discover_roombas( + hass: HomeAssistant, host: str | None = None +) -> list[RoombaInfo]: + discovered_hosts: set[str] = set() + devices: list[RoombaInfo] = [] discover_lock = hass.data.setdefault(ROOMBA_DISCOVERY_LOCK, asyncio.Lock()) discover_attempts = HOST_ATTEMPTS if host else ALL_ATTEMPTS for attempt in range(discover_attempts + 1): async with discover_lock: discovery = _async_get_roomba_discovery() + discovered: set[RoombaInfo] = set() try: if host: device = await hass.async_add_executor_job(discovery.get, host) - discovered = [device] if device else [] + if device: + discovered.add(device) else: discovered = await hass.async_add_executor_job(discovery.get_all) except OSError: