Add emeter support for strip sockets (#203)

* Add support for plugs with emeters.

* Tweaks for emeter

* black

* tweaks

* tweaks

* more tweaks

* dry

* flake8

* flake8

* legacy typing

* Update kasa/smartstrip.py

Co-authored-by: Teemu R. <tpr@iki.fi>

* reduce

* remove useless delegation

* tweaks

* tweaks

* dry

* tweak

* tweak

* tweak

* tweak

* update tests

* wrap

* preen

* prune

* prune

* prune

* guard

* adjust

* robust

* prune

* prune

* reduce dict lookups by 1

* Update kasa/smartstrip.py

Co-authored-by: Teemu R. <tpr@iki.fi>

* delete utils

* isort

Co-authored-by: Brendan Burns <brendan.d.burns@gmail.com>
Co-authored-by: Teemu R. <tpr@iki.fi>

kasa/smartstrip.py

@@ -6,6 +6,7 @@ from typing import Any, DefaultDict, Dict, Optional
 
 from kasa.smartdevice import (
     DeviceType,
+    EmeterStatus,
     SmartDevice,
     SmartDeviceException,
     requires_update,
@@ -15,6 +16,15 @@ from kasa.smartplug import SmartPlug
 _LOGGER = logging.getLogger(__name__)
 
 
+def merge_sums(dicts):
+    """Merge the sum of dicts."""
+    total_dict: DefaultDict[int, float] = defaultdict(lambda: 0.0)
+    for sum_dict in dicts:
+        for day, value in sum_dict.items():
+            total_dict[day] += value
+    return total_dict
+
+
 class SmartStrip(SmartDevice):
     """Representation of a TP-Link Smart Power Strip.
 
@@ -75,11 +85,7 @@ class SmartStrip(SmartDevice):
     @requires_update
     def is_on(self) -> bool:
         """Return if any of the outlets are on."""
-        for plug in self.children:
-            is_on = plug.is_on
-            if is_on:
-                return True
-        return False
+        return any(plug.is_on for plug in self.children)
 
     async def update(self):
         """Update some of the attributes.
@@ -97,6 +103,10 @@ class SmartStrip(SmartDevice):
                     SmartStripPlug(self.host, parent=self, child_id=child["id"])
                 )
 
+        if self.has_emeter:
+            for plug in self.children:
+                await plug.update()
+
     async def turn_on(self, **kwargs):
         """Turn the strip on."""
         await self._query_helper("system", "set_relay_state", {"state": 1})
@@ -140,16 +150,16 @@ class SmartStrip(SmartDevice):
 
     async def current_consumption(self) -> float:
         """Get the current power consumption in watts."""
-        consumption = sum(await plug.current_consumption() for plug in self.children)
+        return sum([await plug.current_consumption() for plug in self.children])
 
-        return consumption
-
-    async def set_alias(self, alias: str) -> None:
-        """Set the alias for the strip.
-
-        :param alias: new alias
-        """
-        return await super().set_alias(alias)
+    @requires_update
+    async def get_emeter_realtime(self) -> EmeterStatus:
+        """Retrieve current energy readings."""
+        emeter_rt = await self._async_get_emeter_sum("get_emeter_realtime", {})
+        # Voltage is averaged since each read will result
+        # in a slightly different voltage since they are not atomic
+        emeter_rt["voltage_mv"] = int(emeter_rt["voltage_mv"] / len(self.children))
+        return EmeterStatus(emeter_rt)
 
     @requires_update
     async def get_emeter_daily(
@@ -163,14 +173,9 @@ class SmartStrip(SmartDevice):
         :param kwh: return usage in kWh (default: True)
         :return: mapping of day of month to value
         """
-        emeter_daily: DefaultDict[int, float] = defaultdict(lambda: 0.0)
-        for plug in self.children:
-            plug_emeter_daily = await plug.get_emeter_daily(
-                year=year, month=month, kwh=kwh
-            )
-            for day, value in plug_emeter_daily.items():
-                emeter_daily[day] += value
-        return emeter_daily
+        return await self._async_get_emeter_sum(
+            "get_emeter_daily", {"year": year, "month": month, "kwh": kwh}
+        )
 
     @requires_update
     async def get_emeter_monthly(self, year: int = None, kwh: bool = True) -> Dict:
@@ -179,13 +184,16 @@ class SmartStrip(SmartDevice):
         :param year: year for which to retrieve statistics (default: this year)
         :param kwh: return usage in kWh (default: True)
         """
-        emeter_monthly: DefaultDict[int, float] = defaultdict(lambda: 0.0)
-        for plug in self.children:
-            plug_emeter_monthly = await plug.get_emeter_monthly(year=year, kwh=kwh)
-            for month, value in plug_emeter_monthly:
-                emeter_monthly[month] += value
+        return await self._async_get_emeter_sum(
+            "get_emeter_monthly", {"year": year, "kwh": kwh}
+        )
 
-        return emeter_monthly
+    async def _async_get_emeter_sum(self, func: str, kwargs: Dict[str, Any]) -> Dict:
+        """Retreive emeter stats for a time period from children."""
+        self._verify_emeter()
+        return merge_sums(
+            [await getattr(plug, func)(**kwargs) for plug in self.children]
+        )
 
     @requires_update
     async def erase_emeter_stats(self):
@@ -193,6 +201,28 @@ class SmartStrip(SmartDevice):
         for plug in self.children:
             await plug.erase_emeter_stats()
 
+    @property  # type: ignore
+    @requires_update
+    def emeter_this_month(self) -> Optional[float]:
+        """Return this month's energy consumption in kWh."""
+        return sum([plug.emeter_this_month for plug in self.children])
+
+    @property  # type: ignore
+    @requires_update
+    def emeter_today(self) -> Optional[float]:
+        """Return this month's energy consumption in kWh."""
+        return sum([plug.emeter_today for plug in self.children])
+
+    @property  # type: ignore
+    @requires_update
+    def emeter_realtime(self) -> EmeterStatus:
+        """Return current energy readings."""
+        emeter = merge_sums([plug.emeter_realtime for plug in self.children])
+        # Voltage is averaged since each read will result
+        # in a slightly different voltage since they are not atomic
+        emeter["voltage_mv"] = int(emeter["voltage_mv"] / len(self.children))
+        return EmeterStatus(emeter)
+
 
 class SmartStripPlug(SmartPlug):
     """Representation of a single socket in a power strip.
@@ -214,12 +244,22 @@ class SmartStripPlug(SmartPlug):
         self._device_type = DeviceType.StripSocket
 
     async def update(self):
-        """Override the update to no-op and inform the user."""
-        _LOGGER.warning(
-            "You called update() on a child device, which has no effect."
-            "Call update() on the parent device instead."
+        """Query the device to update the data.
+
+        Needed for properties that are decorated with `requires_update`.
+        """
+        self._last_update = await self.parent.protocol.query(
+            self.host, self._create_emeter_request()
         )
-        return
+
+    def _create_request(
+        self, target: str, cmd: str, arg: Optional[Dict] = None, child_ids=None
+    ):
+        request: Dict[str, Any] = {
+            "context": {"child_ids": [self.child_id]},
+            target: {cmd: arg},
+        }
+        return request
 
     async def _query_helper(
         self, target: str, cmd: str, arg: Optional[Dict] = None, child_ids=None
@@ -245,12 +285,6 @@ class SmartStripPlug(SmartPlug):
         """
         return False
 
-    @property  # type: ignore
-    @requires_update
-    def has_emeter(self) -> bool:
-        """Children have no emeter to my knowledge."""
-        return False
-
     @property  # type: ignore
     @requires_update
     def device_id(self) -> str: