From 001e14b84345f9856fa0cf4bded7f34825adc63d Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 28 Jun 2023 10:21:35 -0500 Subject: [PATCH 1/7] Refactor ESPHome setup logic into a class Avoids all the nonlocals and fixes the C901 --- homeassistant/components/esphome/__init__.py | 283 ++++++++++++------- 1 file changed, 177 insertions(+), 106 deletions(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index afaefe117baca..3c0e5869674ce 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -137,55 +137,64 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: return True -async def async_setup_entry( # noqa: C901 - hass: HomeAssistant, entry: ConfigEntry -) -> bool: - """Set up the esphome component.""" - host = entry.data[CONF_HOST] - port = entry.data[CONF_PORT] - password = entry.data[CONF_PASSWORD] - noise_psk = entry.data.get(CONF_NOISE_PSK) - device_id: str = None # type: ignore[assignment] - - zeroconf_instance = await zeroconf.async_get_instance(hass) - - cli = APIClient( - host, - port, - password, - client_info=f"Home Assistant {ha_version}", - zeroconf_instance=zeroconf_instance, - noise_psk=noise_psk, +class ESPHomeManager: + """Class to manage an ESPHome connection.""" + + __slots__ = ( + "hass", + "host", + "password", + "entry", + "cli", + "device_id", + "domain_data", + "voice_assistant_udp_server", + "reconnect_logic", + "zeroconf_instance", + "entry_data", ) - services_issue = f"service_calls_not_enabled-{entry.unique_id}" - if entry.options.get(CONF_ALLOW_SERVICE_CALLS, DEFAULT_ALLOW_SERVICE_CALLS): - async_delete_issue(hass, DOMAIN, services_issue) - - domain_data = DomainData.get(hass) - entry_data = RuntimeEntryData( - client=cli, - entry_id=entry.entry_id, - store=domain_data.get_or_create_store(hass, entry), - original_options=dict(entry.options), - ) - domain_data.set_entry_data(entry, entry_data) - - async def on_stop(event: Event) -> None: + def __init__( + self, + hass: HomeAssistant, + entry: ConfigEntry, + host: str, + password: str | None, + cli: APIClient, + zeroconf_instance: zeroconf.HaZeroconf, + domain_data: DomainData, + entry_data: RuntimeEntryData, + ) -> None: + """Initialize the esphome manager.""" + self.hass = hass + self.host = host + self.password = password + self.entry = entry + self.cli = cli + self.device_id: str | None = None + self.domain_data = domain_data + self.voice_assistant_udp_server: VoiceAssistantUDPServer | None = None + self.reconnect_logic: ReconnectLogic | None = None + self.zeroconf_instance = zeroconf_instance + self.entry_data = entry_data + + async def on_stop(self, event: Event) -> None: """Cleanup the socket client on HA stop.""" - await _cleanup_instance(hass, entry) - - # Use async_listen instead of async_listen_once so that we don't deregister - # the callback twice when shutting down Home Assistant. - # "Unable to remove unknown listener - # .onetime_listener>" - entry_data.cleanup_callbacks.append( - hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, on_stop) - ) + await _cleanup_instance(self.hass, self.entry) + + @property + def services_issue(self) -> str: + """Return the services issue name for this entry.""" + return f"service_calls_not_enabled-{self.entry.unique_id}" @callback - def async_on_service_call(service: HomeassistantServiceCall) -> None: + def async_on_service_call(self, service: HomeassistantServiceCall) -> None: """Call service when user automation in ESPHome config is triggered.""" + entry_data = self.entry_data + hass = self.hass + host = self.host + entry = self.entry + device_id = self.device_id device_info = entry_data.device_info assert device_info is not None domain, service_name = service.service.split(".", 1) @@ -236,7 +245,7 @@ def async_on_service_call(service: HomeassistantServiceCall) -> None: async_create_issue( hass, DOMAIN, - services_issue, + self.services_issue, is_fixable=False, severity=IssueSeverity.WARNING, translation_key="service_calls_not_allowed", @@ -256,7 +265,7 @@ def async_on_service_call(service: HomeassistantServiceCall) -> None: ) async def _send_home_assistant_state( - entity_id: str, attribute: str | None, state: State | None + self, entity_id: str, attribute: str | None, state: State | None ) -> None: """Forward Home Assistant states to ESPHome.""" if state is None or (attribute and attribute not in state.attributes): @@ -271,13 +280,15 @@ async def _send_home_assistant_state( else: send_state = attr_val - await cli.send_home_assistant_state(entity_id, attribute, str(send_state)) + await self.cli.send_home_assistant_state(entity_id, attribute, str(send_state)) @callback def async_on_state_subscription( - entity_id: str, attribute: str | None = None + self, entity_id: str, attribute: str | None = None ) -> None: """Subscribe and forward states for requested entities.""" + hass = self.hass + entry_data = self.entry_data async def send_home_assistant_state_event(event: Event) -> None: """Forward Home Assistant states updates to ESPHome.""" @@ -303,7 +314,7 @@ async def send_home_assistant_state_event(event: Event) -> None: ): return - await _send_home_assistant_state( + await self._send_home_assistant_state( event.data["entity_id"], attribute, event.data.get("new_state") ) @@ -314,59 +325,64 @@ async def send_home_assistant_state_event(event: Event) -> None: # Send initial state hass.async_create_task( - _send_home_assistant_state(entity_id, attribute, hass.states.get(entity_id)) + self._send_home_assistant_state( + entity_id, attribute, hass.states.get(entity_id) + ) ) - voice_assistant_udp_server: VoiceAssistantUDPServer | None = None - def _handle_pipeline_event( - event_type: VoiceAssistantEventType, data: dict[str, str] | None + self, event_type: VoiceAssistantEventType, data: dict[str, str] | None ) -> None: - cli.send_voice_assistant_event(event_type, data) + self.cli.send_voice_assistant_event(event_type, data) - def _handle_pipeline_finished() -> None: - nonlocal voice_assistant_udp_server + def _handle_pipeline_finished(self) -> None: + self.entry_data.async_set_assist_pipeline_state(False) - entry_data.async_set_assist_pipeline_state(False) + if self.voice_assistant_udp_server is not None: + self.voice_assistant_udp_server.close() + self.voice_assistant_udp_server = None - if voice_assistant_udp_server is not None: - voice_assistant_udp_server.close() - voice_assistant_udp_server = None - - async def _handle_pipeline_start(conversation_id: str, use_vad: bool) -> int | None: + async def _handle_pipeline_start( + self, conversation_id: str, use_vad: bool + ) -> int | None: """Start a voice assistant pipeline.""" - nonlocal voice_assistant_udp_server - - if voice_assistant_udp_server is not None: + if self.voice_assistant_udp_server is not None: return None + hass = self.hass voice_assistant_udp_server = VoiceAssistantUDPServer( - hass, entry_data, _handle_pipeline_event, _handle_pipeline_finished + hass, + self.entry_data, + self._handle_pipeline_event, + self._handle_pipeline_finished, ) port = await voice_assistant_udp_server.start_server() + assert self.device_id is not None, "Device ID must be set" hass.async_create_background_task( voice_assistant_udp_server.run_pipeline( - device_id=device_id, + device_id=self.device_id, conversation_id=conversation_id or None, use_vad=use_vad, ), "esphome.voice_assistant_udp_server.run_pipeline", ) - entry_data.async_set_assist_pipeline_state(True) + self.entry_data.async_set_assist_pipeline_state(True) return port - async def _handle_pipeline_stop() -> None: + async def _handle_pipeline_stop(self) -> None: """Stop a voice assistant pipeline.""" - nonlocal voice_assistant_udp_server + if self.voice_assistant_udp_server is not None: + self.voice_assistant_udp_server.stop() - if voice_assistant_udp_server is not None: - voice_assistant_udp_server.stop() - - async def on_connect() -> None: + async def on_connect(self) -> None: """Subscribe to states and list entities on successful API login.""" - nonlocal device_id + entry = self.entry + entry_data = self.entry_data + reconnect_logic = self.reconnect_logic + hass = self.hass + cli = self.cli try: device_info = await cli.device_info() @@ -389,6 +405,7 @@ async def on_connect() -> None: entry_data.api_version = cli.api_version entry_data.available = True if entry_data.device_info.name: + assert reconnect_logic is not None, "Reconnect logic must be set" reconnect_logic.name = entry_data.device_info.name if device_info.bluetooth_proxy_feature_flags_compat(cli.api_version): @@ -396,37 +413,38 @@ async def on_connect() -> None: await async_connect_scanner(hass, entry, cli, entry_data) ) - device_id = _async_setup_device_registry( - hass, entry, entry_data.device_info - ) + _async_setup_device_registry(hass, entry, entry_data.device_info) entry_data.async_update_device_state(hass) entity_infos, services = await cli.list_entities_services() await entry_data.async_update_static_infos(hass, entry, entity_infos) await _setup_services(hass, entry_data, services) await cli.subscribe_states(entry_data.async_update_state) - await cli.subscribe_service_calls(async_on_service_call) - await cli.subscribe_home_assistant_states(async_on_state_subscription) + await cli.subscribe_service_calls(self.async_on_service_call) + await cli.subscribe_home_assistant_states(self.async_on_state_subscription) if device_info.voice_assistant_version: entry_data.disconnect_callbacks.append( await cli.subscribe_voice_assistant( - _handle_pipeline_start, - _handle_pipeline_stop, + self._handle_pipeline_start, + self._handle_pipeline_stop, ) ) hass.async_create_task(entry_data.async_save_to_store()) except APIConnectionError as err: - _LOGGER.warning("Error getting initial data for %s: %s", host, err) + _LOGGER.warning("Error getting initial data for %s: %s", self.host, err) # Re-connection logic will trigger after this await cli.disconnect() else: _async_check_firmware_version(hass, device_info, entry_data.api_version) - _async_check_using_api_password(hass, device_info, bool(password)) + _async_check_using_api_password(hass, device_info, bool(self.password)) - async def on_disconnect(expected_disconnect: bool) -> None: + async def on_disconnect(self, expected_disconnect: bool) -> None: """Run disconnect callbacks on API disconnect.""" + entry_data = self.entry_data + hass = self.hass + host = self.host name = entry_data.device_info.name if entry_data.device_info else host _LOGGER.debug( "%s: %s disconnected (expected=%s), running disconnected callbacks", @@ -453,7 +471,7 @@ async def on_disconnect(expected_disconnect: bool) -> None: # will be cleared anyway. entry_data.async_update_device_state(hass) - async def on_connect_error(err: Exception) -> None: + async def on_connect_error(self, err: Exception) -> None: """Start reauth flow if appropriate connect error type.""" if isinstance( err, @@ -463,32 +481,85 @@ async def on_connect_error(err: Exception) -> None: InvalidAuthAPIError, ), ): - entry.async_start_reauth(hass) + self.entry.async_start_reauth(self.hass) + + async def async_start(self) -> None: + """Start the esphome component coordinator.""" + hass = self.hass + entry = self.entry + entry_data = self.entry_data + + if entry.options.get(CONF_ALLOW_SERVICE_CALLS, DEFAULT_ALLOW_SERVICE_CALLS): + async_delete_issue(hass, DOMAIN, self.services_issue) + + # Use async_listen instead of async_listen_once so that we don't deregister + # the callback twice when shutting down Home Assistant. + # "Unable to remove unknown listener + # .onetime_listener>" + entry_data.cleanup_callbacks.append( + hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, self.on_stop) + ) - reconnect_logic = ReconnectLogic( - client=cli, - on_connect=on_connect, - on_disconnect=on_disconnect, - zeroconf_instance=zeroconf_instance, - name=host, - on_connect_error=on_connect_error, - ) + reconnect_logic = ReconnectLogic( + client=self.cli, + on_connect=self.on_connect, + on_disconnect=self.on_disconnect, + zeroconf_instance=self.zeroconf_instance, + name=self.host, + on_connect_error=self.on_connect_error, + ) + self.reconnect_logic = reconnect_logic - infos, services = await entry_data.async_load_from_store() - await entry_data.async_update_static_infos(hass, entry, infos) - await _setup_services(hass, entry_data, services) + infos, services = await entry_data.async_load_from_store() + await entry_data.async_update_static_infos(hass, entry, infos) + await _setup_services(hass, entry_data, services) - if entry_data.device_info is not None and entry_data.device_info.name: - reconnect_logic.name = entry_data.device_info.name - if entry.unique_id is None: - hass.config_entries.async_update_entry( - entry, unique_id=format_mac(entry_data.device_info.mac_address) - ) + if entry_data.device_info is not None and entry_data.device_info.name: + reconnect_logic.name = entry_data.device_info.name + if entry.unique_id is None: + hass.config_entries.async_update_entry( + entry, unique_id=format_mac(entry_data.device_info.mac_address) + ) + + await reconnect_logic.start() + entry_data.cleanup_callbacks.append(reconnect_logic.stop_callback) + + entry.async_on_unload( + entry.add_update_listener(entry_data.async_update_listener) + ) + + +async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + """Set up the esphome component.""" + host = entry.data[CONF_HOST] + port = entry.data[CONF_PORT] + password = entry.data[CONF_PASSWORD] + noise_psk = entry.data.get(CONF_NOISE_PSK) + + zeroconf_instance = await zeroconf.async_get_instance(hass) - await reconnect_logic.start() - entry_data.cleanup_callbacks.append(reconnect_logic.stop_callback) + cli = APIClient( + host, + port, + password, + client_info=f"Home Assistant {ha_version}", + zeroconf_instance=zeroconf_instance, + noise_psk=noise_psk, + ) + + domain_data = DomainData.get(hass) + entry_data = RuntimeEntryData( + client=cli, + entry_id=entry.entry_id, + store=domain_data.get_or_create_store(hass, entry), + original_options=dict(entry.options), + ) + domain_data.set_entry_data(entry, entry_data) - entry.async_on_unload(entry.add_update_listener(entry_data.async_update_listener)) + manager = ESPHomeManager( + hass, entry, host, password, cli, zeroconf_instance, domain_data, entry_data + ) + await manager.async_start() return True From ab740d133a9ef5493ef0e014788de977946d5992 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 28 Jun 2023 10:28:45 -0500 Subject: [PATCH 2/7] cleanup --- homeassistant/components/esphome/__init__.py | 49 +++++++++----------- 1 file changed, 22 insertions(+), 27 deletions(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 3c0e5869674ce..69ea2ae9ccb3a 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -190,13 +190,7 @@ def services_issue(self) -> str: @callback def async_on_service_call(self, service: HomeassistantServiceCall) -> None: """Call service when user automation in ESPHome config is triggered.""" - entry_data = self.entry_data hass = self.hass - host = self.host - entry = self.entry - device_id = self.device_id - device_info = entry_data.device_info - assert device_info is not None domain, service_name = service.service.split(".", 1) service_data = service.data @@ -210,15 +204,16 @@ def async_on_service_call(self, service: HomeassistantServiceCall) -> None: template.render_complex(data_template, service.variables) ) except TemplateError as ex: - _LOGGER.error("Error rendering data template for %s: %s", host, ex) + _LOGGER.error("Error rendering data template for %s: %s", self.host, ex) return if service.is_event: + device_id = self.device_id # ESPHome uses service call packet for both events and service calls # Ensure the user can only send events of form 'esphome.xyz' if domain != "esphome": _LOGGER.error( - "Can only generate events under esphome domain! (%s)", host + "Can only generate events under esphome domain! (%s)", self.host ) return @@ -235,13 +230,17 @@ def async_on_service_call(self, service: HomeassistantServiceCall) -> None: **service_data, }, ) - elif entry.options.get(CONF_ALLOW_SERVICE_CALLS, DEFAULT_ALLOW_SERVICE_CALLS): + elif self.entry.options.get( + CONF_ALLOW_SERVICE_CALLS, DEFAULT_ALLOW_SERVICE_CALLS + ): hass.async_create_task( hass.services.async_call( domain, service_name, service_data, blocking=True ) ) else: + device_info = self.entry_data.device_info + assert device_info is not None async_create_issue( hass, DOMAIN, @@ -288,40 +287,36 @@ def async_on_state_subscription( ) -> None: """Subscribe and forward states for requested entities.""" hass = self.hass - entry_data = self.entry_data async def send_home_assistant_state_event(event: Event) -> None: """Forward Home Assistant states updates to ESPHome.""" - + event_data = event.data # Only communicate changes to the state or attribute tracked - if event.data.get("new_state") is None or ( - event.data.get("old_state") is not None - and "new_state" in event.data + if (new_state := event_data.get("new_state")) is None or ( + (old_state := event_data.get("old_state")) is not None + and "new_state" in event_data and ( - ( - not attribute - and event.data["old_state"].state - == event.data["new_state"].state - ) + (not attribute and old_state.state == new_state.state) or ( attribute - and attribute in event.data["old_state"].attributes - and attribute in event.data["new_state"].attributes - and event.data["old_state"].attributes[attribute] - == event.data["new_state"].attributes[attribute] + and attribute in old_state.attributes + and attribute in new_state.attributes + and old_state.attributes[attribute] + == new_state.attributes[attribute] ) ) ): return await self._send_home_assistant_state( - event.data["entity_id"], attribute, event.data.get("new_state") + event.data["entity_id"], attribute, new_state ) - unsub = async_track_state_change_event( - hass, [entity_id], send_home_assistant_state_event + self.entry_data.disconnect_callbacks.append( + async_track_state_change_event( + hass, [entity_id], send_home_assistant_state_event + ) ) - entry_data.disconnect_callbacks.append(unsub) # Send initial state hass.async_create_task( From 83e06ced6d696da17ad5a46d143c140347e3342f Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 28 Jun 2023 10:30:30 -0500 Subject: [PATCH 3/7] touch ups --- homeassistant/components/esphome/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 69ea2ae9ccb3a..582eeea1a2d2c 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -291,9 +291,11 @@ def async_on_state_subscription( async def send_home_assistant_state_event(event: Event) -> None: """Forward Home Assistant states updates to ESPHome.""" event_data = event.data + new_state = event_data.get("new_state") + old_state = event_data.get("old_state") # Only communicate changes to the state or attribute tracked - if (new_state := event_data.get("new_state")) is None or ( - (old_state := event_data.get("old_state")) is not None + if new_state is None or ( + old_state is not None and "new_state" in event_data and ( (not attribute and old_state.state == new_state.state) From 33dd4f2b2c961cd7e59c6326281c37cab14dc6a0 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 28 Jun 2023 10:31:27 -0500 Subject: [PATCH 4/7] touch ups --- homeassistant/components/esphome/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 582eeea1a2d2c..9ffecf1370b71 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -291,8 +291,8 @@ def async_on_state_subscription( async def send_home_assistant_state_event(event: Event) -> None: """Forward Home Assistant states updates to ESPHome.""" event_data = event.data - new_state = event_data.get("new_state") - old_state = event_data.get("old_state") + new_state: State | None = event_data.get("new_state") + old_state: State | None = event_data.get("old_state") # Only communicate changes to the state or attribute tracked if new_state is None or ( old_state is not None From d8e1a356e99468434ea3e7b5193d0eb0305027ba Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 28 Jun 2023 10:31:46 -0500 Subject: [PATCH 5/7] touch ups --- homeassistant/components/esphome/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 9ffecf1370b71..eb170f84cfb5e 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -296,7 +296,7 @@ async def send_home_assistant_state_event(event: Event) -> None: # Only communicate changes to the state or attribute tracked if new_state is None or ( old_state is not None - and "new_state" in event_data + and new_state and ( (not attribute and old_state.state == new_state.state) or ( From 56a53235145f55204357d2ee3e81219b4c5181f0 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 28 Jun 2023 10:36:13 -0500 Subject: [PATCH 6/7] make easier to read --- homeassistant/components/esphome/__init__.py | 21 ++++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index eb170f84cfb5e..d8c11adcc2993 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -293,20 +293,15 @@ async def send_home_assistant_state_event(event: Event) -> None: event_data = event.data new_state: State | None = event_data.get("new_state") old_state: State | None = event_data.get("old_state") + + if new_state is None or old_state is None: + return + # Only communicate changes to the state or attribute tracked - if new_state is None or ( - old_state is not None - and new_state - and ( - (not attribute and old_state.state == new_state.state) - or ( - attribute - and attribute in old_state.attributes - and attribute in new_state.attributes - and old_state.attributes[attribute] - == new_state.attributes[attribute] - ) - ) + if (not attribute and old_state.state == new_state.state) or ( + attribute + and old_state.attributes.get(attribute) + == new_state.attributes.get(attribute) ): return From c73a58d60b00aabfa46244af325df960ce2e83d4 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 28 Jun 2023 10:38:57 -0500 Subject: [PATCH 7/7] stale --- homeassistant/components/esphome/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index d8c11adcc2993..271b0b9aa1655 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -476,7 +476,7 @@ async def on_connect_error(self, err: Exception) -> None: self.entry.async_start_reauth(self.hass) async def async_start(self) -> None: - """Start the esphome component coordinator.""" + """Start the esphome connection manager.""" hass = self.hass entry = self.entry entry_data = self.entry_data