44
55from abc import ABC , abstractmethod
66import asyncio
7+ from enum import StrEnum
78import logging
89from typing import Any
910
2324 ConfigEntryBaseFlow ,
2425 ConfigFlow ,
2526 ConfigFlowResult ,
27+ FlowType ,
2628 OptionsFlow ,
2729)
2830from homeassistant .core import callback
5052STEP_PICK_FIRMWARE_ZIGBEE = "pick_firmware_zigbee"
5153
5254
55+ class PickedFirmwareType (StrEnum ):
56+ """Firmware types that can be picked."""
57+
58+ THREAD = "thread"
59+ ZIGBEE = "zigbee"
60+
61+
62+ class ZigbeeIntegration (StrEnum ):
63+ """Zigbee integrations that can be picked."""
64+
65+ OTHER = "other"
66+ ZHA = "zha"
67+
68+
5369class BaseFirmwareInstallFlow (ConfigEntryBaseFlow , ABC ):
5470 """Base flow to install firmware."""
5571
72+ ZIGBEE_BAUDRATE = 115200 # Default, subclasses may override
5673 _failed_addon_name : str
5774 _failed_addon_reason : str
75+ _picked_firmware_type : PickedFirmwareType
5876
5977 def __init__ (self , * args : Any , ** kwargs : Any ) -> None :
6078 """Instantiate base flow."""
@@ -63,6 +81,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
6381 self ._probed_firmware_info : FirmwareInfo | None = None
6482 self ._device : str | None = None # To be set in a subclass
6583 self ._hardware_name : str = "unknown" # To be set in a subclass
84+ self ._zigbee_integration = ZigbeeIntegration .ZHA
6685
6786 self .addon_install_task : asyncio .Task | None = None
6887 self .addon_start_task : asyncio .Task | None = None
@@ -281,17 +300,79 @@ async def async_step_firmware_install_failed(
281300 },
282301 )
283302
284- async def async_step_pick_firmware_zigbee (
303+ async def async_step_zigbee_installation_type (
285304 self , user_input : dict [str , Any ] | None = None
286305 ) -> ConfigFlowResult :
287- """Pick Zigbee firmware."""
306+ """Handle the installation type step."""
307+ return self .async_show_menu (
308+ step_id = "zigbee_installation_type" ,
309+ menu_options = [
310+ "zigbee_intent_recommended" ,
311+ "zigbee_intent_custom" ,
312+ ],
313+ )
314+
315+ async def async_step_zigbee_intent_recommended (
316+ self , user_input : dict [str , Any ] | None = None
317+ ) -> ConfigFlowResult :
318+ """Select recommended installation type."""
319+ self ._zigbee_integration = ZigbeeIntegration .ZHA
320+ return await self ._async_continue_picked_firmware ()
321+
322+ async def async_step_zigbee_intent_custom (
323+ self , user_input : dict [str , Any ] | None = None
324+ ) -> ConfigFlowResult :
325+ """Select custom installation type."""
326+ return await self .async_step_zigbee_integration ()
327+
328+ async def async_step_zigbee_integration (
329+ self , user_input : dict [str , Any ] | None = None
330+ ) -> ConfigFlowResult :
331+ """Select Zigbee integration."""
332+ return self .async_show_menu (
333+ step_id = "zigbee_integration" ,
334+ menu_options = [
335+ "zigbee_integration_zha" ,
336+ "zigbee_integration_other" ,
337+ ],
338+ )
339+
340+ async def async_step_zigbee_integration_zha (
341+ self , user_input : dict [str , Any ] | None = None
342+ ) -> ConfigFlowResult :
343+ """Select ZHA integration."""
344+ self ._zigbee_integration = ZigbeeIntegration .ZHA
345+ return await self ._async_continue_picked_firmware ()
346+
347+ async def async_step_zigbee_integration_other (
348+ self , user_input : dict [str , Any ] | None = None
349+ ) -> ConfigFlowResult :
350+ """Select other Zigbee integration."""
351+ self ._zigbee_integration = ZigbeeIntegration .OTHER
352+ return await self ._async_continue_picked_firmware ()
353+
354+ async def _async_continue_picked_firmware (self ) -> ConfigFlowResult :
355+ """Continue to the picked firmware step."""
288356 if not await self ._probe_firmware_info ():
289357 return self .async_abort (
290358 reason = "unsupported_firmware" ,
291359 description_placeholders = self ._get_translation_placeholders (),
292360 )
293361
294- return await self .async_step_install_zigbee_firmware ()
362+ if self ._picked_firmware_type == PickedFirmwareType .ZIGBEE :
363+ return await self .async_step_install_zigbee_firmware ()
364+
365+ if result := await self ._ensure_thread_addon_setup ():
366+ return result
367+
368+ return await self .async_step_install_thread_firmware ()
369+
370+ async def async_step_pick_firmware_zigbee (
371+ self , user_input : dict [str , Any ] | None = None
372+ ) -> ConfigFlowResult :
373+ """Pick Zigbee firmware."""
374+ self ._picked_firmware_type = PickedFirmwareType .ZIGBEE
375+ return await self .async_step_zigbee_installation_type ()
295376
296377 async def async_step_install_zigbee_firmware (
297378 self , user_input : dict [str , Any ] | None = None
@@ -317,42 +398,43 @@ async def async_step_pre_confirm_zigbee(
317398 """Pre-confirm Zigbee setup."""
318399
319400 # This step is necessary to prevent `user_input` from being passed through
320- return await self .async_step_confirm_zigbee ()
401+ return await self .async_step_continue_zigbee ()
321402
322- async def async_step_confirm_zigbee (
403+ async def async_step_continue_zigbee (
323404 self , user_input : dict [str , Any ] | None = None
324405 ) -> ConfigFlowResult :
325- """Confirm Zigbee setup."""
406+ """Continue Zigbee setup."""
326407 assert self ._device is not None
327408 assert self ._hardware_name is not None
328409
329- if user_input is None :
330- return self .async_show_form (
331- step_id = "confirm_zigbee" ,
332- description_placeholders = self ._get_translation_placeholders (),
333- )
334-
335410 if not await self ._probe_firmware_info (probe_methods = (ApplicationType .EZSP ,)):
336411 return self .async_abort (
337412 reason = "unsupported_firmware" ,
338413 description_placeholders = self ._get_translation_placeholders (),
339414 )
340415
341- await self .hass .config_entries .flow .async_init (
416+ if self ._zigbee_integration == ZigbeeIntegration .OTHER :
417+ return self ._async_flow_finished ()
418+
419+ result = await self .hass .config_entries .flow .async_init (
342420 ZHA_DOMAIN ,
343421 context = {"source" : "hardware" },
344422 data = {
345423 "name" : self ._hardware_name ,
346424 "port" : {
347425 "path" : self ._device ,
348- "baudrate" : 115200 ,
426+ "baudrate" : self . ZIGBEE_BAUDRATE ,
349427 "flow_control" : "hardware" ,
350428 },
351429 "radio_type" : "ezsp" ,
352430 },
353431 )
432+ return self ._continue_zha_flow (result )
354433
355- return self ._async_flow_finished ()
434+ @callback
435+ def _continue_zha_flow (self , zha_result : ConfigFlowResult ) -> ConfigFlowResult :
436+ """Continue the ZHA flow."""
437+ raise NotImplementedError
356438
357439 async def _ensure_thread_addon_setup (self ) -> ConfigFlowResult | None :
358440 """Ensure the OTBR addon is set up and not running."""
@@ -391,16 +473,8 @@ async def async_step_pick_firmware_thread(
391473 self , user_input : dict [str , Any ] | None = None
392474 ) -> ConfigFlowResult :
393475 """Pick Thread firmware."""
394- if not await self ._probe_firmware_info ():
395- return self .async_abort (
396- reason = "unsupported_firmware" ,
397- description_placeholders = self ._get_translation_placeholders (),
398- )
399-
400- if result := await self ._ensure_thread_addon_setup ():
401- return result
402-
403- return await self .async_step_install_thread_firmware ()
476+ self ._picked_firmware_type = PickedFirmwareType .THREAD
477+ return await self ._async_continue_picked_firmware ()
404478
405479 async def async_step_install_thread_firmware (
406480 self , user_input : dict [str , Any ] | None = None
@@ -572,6 +646,21 @@ async def async_step_confirm(
572646
573647 return await self .async_step_pick_firmware ()
574648
649+ @callback
650+ def _continue_zha_flow (self , zha_result : ConfigFlowResult ) -> ConfigFlowResult :
651+ """Continue the ZHA flow."""
652+ next_flow_id = zha_result ["flow_id" ]
653+
654+ result = self ._async_flow_finished ()
655+ return (
656+ self .async_create_entry (
657+ title = result ["title" ] or self ._hardware_name ,
658+ data = result ["data" ],
659+ next_flow = (FlowType .CONFIG_FLOW , next_flow_id ),
660+ )
661+ | result # update all items with the child result
662+ )
663+
575664
576665class BaseFirmwareOptionsFlow (BaseFirmwareInstallFlow , OptionsFlow ):
577666 """Zigbee and Thread options flow handlers."""
@@ -629,3 +718,10 @@ async def async_step_pick_firmware_thread(
629718 )
630719
631720 return await super ().async_step_pick_firmware_thread (user_input )
721+
722+ @callback
723+ def _continue_zha_flow (self , zha_result : ConfigFlowResult ) -> ConfigFlowResult :
724+ """Continue the ZHA flow."""
725+ # The options flow cannot return a next_flow yet, so we just finish here.
726+ # The options flow should be changed to a reconfigure flow.
727+ return self ._async_flow_finished ()
0 commit comments