Refactor basic switch structure

homeassistant/components/switch/__init__.py

@@ -7,6 +7,7 @@ import logging
 from datetime import timedelta
 
 from homeassistant.helpers.entity_component import EntityComponent
+from homeassistant.helpers.entity import ToggleEntity
 
 from homeassistant.const import (
     STATE_ON, SERVICE_TURN_ON, SERVICE_TURN_OFF, ATTR_ENTITY_ID)
@@ -33,6 +34,11 @@ DISCOVERY_PLATFORMS = {
     isy994.DISCOVER_SWITCHES: 'isy994',
 }
 
+PROP_TO_ATTR = {
+    'current_power_mwh': ATTR_CURRENT_POWER_MWH,
+    'today_power_mw': ATTR_TODAY_MWH,
+}
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -74,10 +80,48 @@ def setup(hass, config):
             else:
                 switch.turn_off()
 
-            switch.update_ha_state(True)
+            if switch.should_poll:
+                switch.update_ha_state(True)
 
     hass.services.register(DOMAIN, SERVICE_TURN_OFF, handle_switch_service)
 
     hass.services.register(DOMAIN, SERVICE_TURN_ON, handle_switch_service)
 
     return True
+
+
+class SwitchDevice(ToggleEntity):
+    """ Represents a switch within Home Assistant. """
+    # pylint: disable=no-self-use
+
+    @property
+    def current_power_mwh(self):
+        """ Current power usage in mwh. """
+        return None
+
+    @property
+    def today_power_mw(self):
+        """ Today total power usage in mw. """
+        return None
+
+    @property
+    def device_state_attributes(self):
+        """ Returns device specific state attributes. """
+        return None
+
+    @property
+    def state_attributes(self):
+        """ Returns optional state attributes. """
+        data = {}
+
+        for prop, attr in PROP_TO_ATTR.items():
+            value = getattr(self, prop)
+            if value:
+                data[attr] = value
+
+        device_attr = self.device_state_attributes
+
+        if device_attr is not None:
+            data.update(device_attr)
+
+        return data

homeassistant/components/switch/command_switch.py

@@ -6,8 +6,7 @@ homeassistant.components.switch.command_switch
 Allows to configure custom shell commands to turn a switch on/off.
 """
 import logging
-from homeassistant.helpers.entity import ToggleEntity
-from homeassistant.const import STATE_ON, STATE_OFF, DEVICE_DEFAULT_NAME
+from homeassistant.components.switch import SwitchDevice
 import subprocess
 
 _LOGGER = logging.getLogger(__name__)
@@ -30,11 +29,11 @@ def setup_platform(hass, config, add_devices_callback, discovery_info=None):
     add_devices_callback(devices)
 
 
-class CommandSwitch(ToggleEntity):
+class CommandSwitch(SwitchDevice):
     """ Represents a switch that can be togggled using shell commands """
     def __init__(self, name, command_on, command_off):
-        self._name = name or DEVICE_DEFAULT_NAME
-        self._state = STATE_OFF
+        self._name = name
+        self._state = False
         self._command_on = command_on
         self._command_off = command_off
 
@@ -60,22 +59,19 @@ class CommandSwitch(ToggleEntity):
         """ The name of the switch """
         return self._name
 
-    @property
-    def state(self):
-        """ Returns the state of the switch. """
-        return self._state
-
     @property
     def is_on(self):
         """ True if device is on. """
-        return self._state == STATE_ON
+        return self._state
 
     def turn_on(self, **kwargs):
         """ Turn the device on. """
         if CommandSwitch._switch(self._command_on):
-            self._state = STATE_ON
+            self._state = True
+        self.update_ha_state()
 
     def turn_off(self, **kwargs):
         """ Turn the device off. """
         if CommandSwitch._switch(self._command_off):
-            self._state = STATE_OFF
+            self._state = False
+        self.update_ha_state()

homeassistant/components/switch/demo.py

@@ -5,20 +5,20 @@ homeassistant.components.switch.demo
 Demo platform that has two fake switches.
 
 """
-from homeassistant.helpers.entity import ToggleEntity
-from homeassistant.const import STATE_ON, STATE_OFF, DEVICE_DEFAULT_NAME
+from homeassistant.components.switch import SwitchDevice
+from homeassistant.const import DEVICE_DEFAULT_NAME
 
 
 # pylint: disable=unused-argument
 def setup_platform(hass, config, add_devices_callback, discovery_info=None):
     """ Find and return demo switches. """
     add_devices_callback([
-        DemoSwitch('Ceiling', STATE_ON),
-        DemoSwitch('AC', STATE_OFF)
+        DemoSwitch('Ceiling', True),
+        DemoSwitch('AC', False)
     ])
 
 
-class DemoSwitch(ToggleEntity):
+class DemoSwitch(SwitchDevice):
     """ Provides a demo switch. """
     def __init__(self, name, state):
         self._name = name or DEVICE_DEFAULT_NAME
@@ -35,19 +35,27 @@ class DemoSwitch(ToggleEntity):
         return self._name
 
     @property
-    def state(self):
-        """ Returns the state of the device if any. """
-        return self._state
+    def current_power_mwh(self):
+        """ Current power usage in mwh. """
+        if self._state:
+            return 100
+
+    @property
+    def today_power_mw(self):
+        """ Today total power usage in mw. """
+        return 1500
 
     @property
     def is_on(self):
         """ True if device is on. """
-        return self._state == STATE_ON
+        return self._state
 
     def turn_on(self, **kwargs):
         """ Turn the device on. """
-        self._state = STATE_ON
+        self._state = True
+        self.update_ha_state()
 
     def turn_off(self, **kwargs):
         """ Turn the device off. """
-        self._state = STATE_OFF
+        self._state = False
+        self.update_ha_state()

homeassistant/components/switch/wemo.py

@@ -6,9 +6,7 @@ Support for WeMo switches.
 """
 import logging
 
-from homeassistant.helpers.entity import ToggleEntity
-from homeassistant.components.switch import (
-    ATTR_TODAY_MWH, ATTR_CURRENT_POWER_MWH)
+from homeassistant.components.switch import SwitchDevice
 
 
 # pylint: disable=unused-argument
@@ -43,10 +41,11 @@ def setup_platform(hass, config, add_devices_callback, discovery_info=None):
          if isinstance(switch, pywemo.Switch)])
 
 
-class WemoSwitch(ToggleEntity):
+class WemoSwitch(SwitchDevice):
     """ Represents a WeMo switch within Home Assistant. """
     def __init__(self, wemo):
         self.wemo = wemo
+        self.insight_params = None
 
     @property
     def unique_id(self):
@@ -59,15 +58,16 @@ class WemoSwitch(ToggleEntity):
         return self.wemo.name
 
     @property
-    def state_attributes(self):
-        """ Returns optional state attributes. """
-        if self.wemo.model.startswith('Belkin Insight'):
-            cur_info = self.wemo.insight_params
+    def current_power_mwh(self):
+        """ Current power usage in mwh. """
+        if self.insight_params:
+            return self.insight_params['currentpower']
 
-            return {
-                ATTR_CURRENT_POWER_MWH: cur_info['currentpower'],
-                ATTR_TODAY_MWH: cur_info['todaymw']
-            }
+    @property
+    def today_power_mw(self):
+        """ Today total power usage in mw. """
+        if self.insight_params:
+            return self.insight_params['todaymw']
 
     @property
     def is_on(self):
@@ -85,3 +85,5 @@ class WemoSwitch(ToggleEntity):
     def update(self):
         """ Update WeMo state. """
         self.wemo.get_state(True)
+        if self.wemo.model.startswith('Belkin Insight'):
+            self.insight_params = self.wemo.insight_params