From c5d8792c3487e9b418b1e7d623bf59e7dbddd6b7 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Fri, 31 Dec 2021 12:47:16 +0100
Subject: [PATCH] Cleanup switch light (#63127)

---
 homeassistant/components/switch/light.py | 79 +++++++++---------------
 1 file changed, 28 insertions(+), 51 deletions(-)

diff --git a/homeassistant/components/switch/light.py b/homeassistant/components/switch/light.py
index 9ca5df135bd..2a83a767cc0 100644
--- a/homeassistant/components/switch/light.py
+++ b/homeassistant/components/switch/light.py
@@ -1,7 +1,7 @@
 """Light support for switch entities."""
 from __future__ import annotations
 
-from typing import Any, cast
+from typing import Any
 
 import voluptuous as vol
 
@@ -18,14 +18,13 @@ from homeassistant.const import (
     STATE_ON,
     STATE_UNAVAILABLE,
 )
-from homeassistant.core import HomeAssistant, State, callback
+from homeassistant.core import Event, HomeAssistant, callback
+from homeassistant.helpers import entity_registry as er
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_state_change_event
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
-# mypy: allow-untyped-calls, allow-untyped-defs, no-check-untyped-defs
-
 DEFAULT_NAME = "Light Switch"
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
@@ -43,15 +42,14 @@ async def async_setup_platform(
     discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
     """Initialize Light Switch platform."""
-
-    registry = await hass.helpers.entity_registry.async_get_registry()
+    registry = er.async_get(hass)
     wrapped_switch = registry.async_get(config[CONF_ENTITY_ID])
     unique_id = wrapped_switch.unique_id if wrapped_switch else None
 
     async_add_entities(
         [
             LightSwitch(
-                cast(str, config.get(CONF_NAME)),
+                config[CONF_NAME],
                 config[CONF_ENTITY_ID],
                 unique_id,
             )
@@ -62,74 +60,50 @@ async def async_setup_platform(
 class LightSwitch(LightEntity):
     """Represents a Switch as a Light."""
 
-    def __init__(self, name: str, switch_entity_id: str, unique_id: str) -> None:
+    _attr_color_mode = COLOR_MODE_ONOFF
+    _attr_should_poll = False
+    _attr_supported_color_modes = {COLOR_MODE_ONOFF}
+
+    def __init__(self, name: str, switch_entity_id: str, unique_id: str | None) -> None:
         """Initialize Light Switch."""
-        self._name = name
+        self._attr_name = name
+        self._attr_unique_id = unique_id
         self._switch_entity_id = switch_entity_id
-        self._unique_id = unique_id
-        self._switch_state: State | None = None
-        self._attr_color_mode = COLOR_MODE_ONOFF
-        self._attr_supported_color_modes = {COLOR_MODE_ONOFF}
 
-    @property
-    def name(self) -> str:
-        """Return the name of the entity."""
-        return self._name
-
-    @property
-    def is_on(self) -> bool:
-        """Return true if light switch is on."""
-        assert self._switch_state is not None
-        return self._switch_state.state == STATE_ON
-
-    @property
-    def available(self) -> bool:
-        """Return true if light switch is on."""
-        return (
-            self._switch_state is not None
-            and self._switch_state.state != STATE_UNAVAILABLE
-        )
-
-    @property
-    def should_poll(self) -> bool:
-        """No polling needed for a light switch."""
-        return False
-
-    @property
-    def unique_id(self):
-        """Return the unique id of the light switch."""
-        return self._unique_id
-
-    async def async_turn_on(self, **kwargs):
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Forward the turn_on command to the switch in this light switch."""
-        data = {ATTR_ENTITY_ID: self._switch_entity_id}
         await self.hass.services.async_call(
             switch.DOMAIN,
             switch.SERVICE_TURN_ON,
-            data,
+            {ATTR_ENTITY_ID: self._switch_entity_id},
             blocking=True,
             context=self._context,
         )
 
-    async def async_turn_off(self, **kwargs):
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Forward the turn_off command to the switch in this light switch."""
-        data = {ATTR_ENTITY_ID: self._switch_entity_id}
         await self.hass.services.async_call(
             switch.DOMAIN,
             switch.SERVICE_TURN_OFF,
-            data,
+            {ATTR_ENTITY_ID: self._switch_entity_id},
             blocking=True,
             context=self._context,
         )
 
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""
-        self._switch_state = self.hass.states.get(self._switch_entity_id)
 
         @callback
-        def async_state_changed_listener(*_: Any) -> None:
+        def async_state_changed_listener(event: Event | None = None) -> None:
             """Handle child updates."""
-            self._switch_state = self.hass.states.get(self._switch_entity_id)
+            if (
+                state := self.hass.states.get(self._switch_entity_id)
+            ) is None or state.state == STATE_UNAVAILABLE:
+                self._attr_available = False
+                return
+
+            self._attr_available = True
+            self._attr_is_on = state.state == STATE_ON
             self.async_write_ha_state()
 
         self.async_on_remove(
@@ -137,3 +111,6 @@ class LightSwitch(LightEntity):
                 self.hass, [self._switch_entity_id], async_state_changed_listener
             )
         )
+
+        # Call once on adding
+        async_state_changed_listener()