python-kasa · rytilahti · Feb 18, 2023 · Nov 7, 2022 · Nov 8, 2022 · Nov 8, 2022
diff --git a/kasa/modules/emeter.py b/kasa/modules/emeter.py
@@ -19,7 +19,7 @@ def emeter_today(self) -> Optional[float]:
         """Return today's energy consumption in kWh."""
         raw_data = self.daily_data
         today = datetime.now().day
-        data = self._emeter_convert_emeter_data(raw_data)
+        data = self._convert_stat_data(raw_data, entry_key="day")
 
         return data.get(today)
 
@@ -28,7 +28,7 @@ def emeter_this_month(self) -> Optional[float]:
         """Return this month's energy consumption in kWh."""
         raw_data = self.monthly_data
         current_month = datetime.now().month
-        data = self._emeter_convert_emeter_data(raw_data)
+        data = self._convert_stat_data(raw_data, entry_key="month")
 
         return data.get(current_month)
 
@@ -43,31 +43,46 @@ async def get_realtime(self):
         """Return real-time statistics."""
         return await self.call("get_realtime")
 
-    async def get_daystat(self, *, year, month, kwh=True):
-        """Return daily stats for the given year & month."""
-        raw_data = await super().get_daystat(year=year, month=month)
-        return self._emeter_convert_emeter_data(raw_data["day_list"], kwh)
+    async def get_daystat(self, *, year=None, month=None, kwh=True) -> Dict:
+        """Return daily stats for the given year & month as a dictionary of {day: energy, ...}."""
+        data = await self.get_raw_daystat(year=year, month=month)
+        data = self._convert_stat_data(data["day_list"], entry_key="day", kwh=kwh)
+        return data
+
+    async def get_monthstat(self, *, year=None, kwh=True) -> Dict:
+        """Return monthly stats for the given year as a dictionary of {month: energy, ...}."""
+        data = await self.get_raw_monthstat(year=year)
+        data = self._convert_stat_data(data["month_list"], entry_key="month", kwh=kwh)
+        return data
 
-    async def get_monthstat(self, *, year, kwh=True):
-        """Return monthly stats for the given year."""
-        raw_data = await super().get_monthstat(year=year)
-        return self._emeter_convert_emeter_data(raw_data["month_list"], kwh)
+    def _convert_stat_data(self, data, entry_key, kwh=True) -> Dict:
+        """Return emeter information keyed with the day/month.
 
-    def _emeter_convert_emeter_data(self, data, kwh=True) -> Dict:
-        """Return emeter information keyed with the day/month.."""
-        response = [EmeterStatus(**x) for x in data]
+        The incoming data is a list of dictionaries::
 
-        if not response:
+            [{'year':      int,
+              'month':     int,
+              'day':       int,     <-- for get_daystat not get_monthstat
+              'energy_wh': int,     <-- for emeter in some versions (wh)
+              'energy':    float    <-- for emeter in other versions (kwh)
+            }, ...]
+
+        :return: a dictionary keyed by day or month with energy as the value.
+        """
+        if not data:
             return {}
 
-        energy_key = "energy_wh"
-        if kwh:
-            energy_key = "energy"
+        scale: float = 1
 
-        entry_key = "month"
-        if "day" in response[0]:
-            entry_key = "day"
+        if "energy_wh" in data[0]:
+            value_key = "energy_wh"
+            if kwh:
+                scale = 1 / 1000
+        else:
+            value_key = "energy"
+            if not kwh:
+                scale = 1000
 
-        data = {entry[entry_key]: entry[energy_key] for entry in response}
+        data = {entry[entry_key]: entry[value_key] * scale for entry in data}
 
         return data
diff --git a/kasa/modules/usage.py b/kasa/modules/usage.py
@@ -1,5 +1,6 @@
 """Implementation of the usage interface."""
 from datetime import datetime
+from typing import Dict
 
 from .module import Module, merge
 
@@ -50,22 +51,54 @@ def usage_this_month(self):
 
         return converted.pop()
 
-    async def get_daystat(self, *, year=None, month=None):
-        """Return daily stats for the given year & month."""
+    async def get_raw_daystat(self, *, year=None, month=None) -> Dict:
+        """Return raw daily stats for the given year & month."""
         if year is None:
             year = datetime.now().year
         if month is None:
             month = datetime.now().month
 
         return await self.call("get_daystat", {"year": year, "month": month})
 
-    async def get_monthstat(self, *, year=None):
-        """Return monthly stats for the given year."""
+    async def get_raw_monthstat(self, *, year=None) -> Dict:
+        """Return raw monthly stats for the given year."""
         if year is None:
             year = datetime.now().year
 
         return await self.call("get_monthstat", {"year": year})
 
+    async def get_daystat(self, *, year=None, month=None) -> Dict:
+        """Return daily stats for the given year & month as a dictionary of {day: time, ...}."""
+        data = await self.get_raw_daystat(year=year, month=month)
+        data = self._convert_stat_data(data["day_list"], entry_key="day")
+        return data
+
+    async def get_monthstat(self, *, year=None) -> Dict:
+        """Return monthly stats for the given year as a dictionary of {month: time, ...}."""
+        data = await self.get_raw_monthstat(year=year)
+        data = self._convert_stat_data(data["month_list"], entry_key="month")
+        return data
+
     async def erase_stats(self):
         """Erase all stats."""
         return await self.call("erase_runtime_stat")
+
+    def _convert_stat_data(self, data, entry_key) -> Dict:
+        """Return usage information keyed with the day/month.
+
+        The incoming data is a list of dictionaries::
+
+               [{'year':      int,
+                 'month':     int,
+                 'day':       int,     <-- for get_daystat not get_monthstat
+                 'time':      int,     <-- for usage (mins)
+               }, ...]
+
+        :return: return a dictionary keyed by day or month with time as the value.
+        """
+        if not data:
+            return {}
+
+        data = {entry[entry_key]: entry["time"] for entry in data}
+
+        return data
diff --git a/kasa/smartdevice.py b/kasa/smartdevice.py
@@ -490,25 +490,6 @@ def emeter_this_month(self) -> Optional[float]:
         self._verify_emeter()
         return self.modules["emeter"].emeter_this_month
 
-    def _emeter_convert_emeter_data(self, data, kwh=True) -> Dict:
-        """Return emeter information keyed with the day/month.."""
-        response = [EmeterStatus(**x) for x in data]
-
-        if not response:
-            return {}
-
-        energy_key = "energy_wh"
-        if kwh:
-            energy_key = "energy"
-
-        entry_key = "month"
-        if "day" in response[0]:
-            entry_key = "day"
-
-        data = {entry[entry_key]: entry[energy_key] for entry in response}
-
-        return data
-
     async def get_emeter_daily(
         self, year: Optional[int] = None, month: Optional[int] = None, kwh: bool = True
     ) -> Dict:

diff --git a/kasa/tests/test_usage.py b/kasa/tests/test_usage.py
@@ -0,0 +1,22 @@
+import pytest
+
+from kasa.modules import Usage
+
+
+def test_usage_convert_stat_data():
+    usage = Usage(None, module="usage")
+
+    test_data = []
+    assert usage._convert_stat_data(test_data, "day") == {}
+
+    test_data = [
+        {"year": 2016, "month": 5, "day": 2, "time": 20},
+        {"year": 2016, "month": 5, "day": 4, "time": 30},
+    ]
+    d = usage._convert_stat_data(test_data, "day")
+    assert len(d) == len(test_data)
+    assert isinstance(d, dict)
+    k, v = d.popitem()
+    assert isinstance(k, int)
+    assert isinstance(v, int)
+    assert k == 4 and v == 30