Home-Assistant
/
core
mirror of https://github.com/home-assistant/core.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add typing to Roomba config flow (#114624)

pull/114668/head
Joost Lekkerkerker 2024-04-02 20:21:55 +02:00 committed by GitHub
parent f85511255c
commit 7cb01f75ae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 55 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

106
homeassistant/components/roomba/config_flow.py
View File

@ -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: