From ff2a6c5e1029a30d01edfd3c60abdbb715f680e1 Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Sat, 23 Mar 2024 16:47:43 +0100
Subject: [PATCH] Add support for feature units

---
 kasa/cli.py                        |  8 ++------
 kasa/feature.py                    |  2 ++
 kasa/smart/modules/energymodule.py | 29 ++++++++++++++---------------
 3 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/kasa/cli.py b/kasa/cli.py
index 41a7759e3..b573b4dd7 100755
--- a/kasa/cli.py
+++ b/kasa/cli.py
@@ -614,12 +614,8 @@ async def state(ctx, dev: Device):
 
     echo("\n\t[bold]== Device-specific information == [/bold]")
     for id_, feature in dev.features.items():
-        echo(f"\t{feature.name} ({id_}): {feature.value}")
-
-    if dev.has_emeter:
-        echo("\n\t[bold]== Current State ==[/bold]")
-        emeter_status = dev.emeter_realtime
-        echo(f"\t{emeter_status}")
+        unit = f" {feature.unit}" if feature.unit else ""
+        echo(f"\t{feature.name} ({id_}): {feature.value}{unit}")
 
     echo("\n\t[bold]== Modules ==[/bold]")
     for module in dev.modules.values():
diff --git a/kasa/feature.py b/kasa/feature.py
index a04e1140a..6add0091a 100644
--- a/kasa/feature.py
+++ b/kasa/feature.py
@@ -36,6 +36,8 @@ class Feature:
     container: Any = None
     #: Icon suggestion
     icon: str | None = None
+    #: Unit, if applicable
+    unit: str | None = None
     #: Type of the feature
     type: FeatureType = FeatureType.Sensor
 
diff --git a/kasa/smart/modules/energymodule.py b/kasa/smart/modules/energymodule.py
index a3e0b4a1c..aedc71aec 100644
--- a/kasa/smart/modules/energymodule.py
+++ b/kasa/smart/modules/energymodule.py
@@ -25,24 +25,27 @@ def __init__(self, device: SmartDevice, module: str):
                 name="Current consumption",
                 attribute_getter="current_power",
                 container=self,
+                unit="W",
             )
-        )  # W or mW?
+        )
         self._add_feature(
             Feature(
                 device,
                 name="Today's consumption",
                 attribute_getter="emeter_today",
                 container=self,
+                unit="Wh",
             )
-        )  # Wh or kWh?
+        )
         self._add_feature(
             Feature(
                 device,
                 name="This month's consumption",
                 attribute_getter="emeter_this_month",
                 container=self,
+                unit="Wh",
             )
-        )  # Wh or kWH?
+        )
 
     def query(self) -> dict:
         """Query to execute during the update cycle."""
@@ -54,9 +57,11 @@ def query(self) -> dict:
         return req
 
     @property
-    def current_power(self):
-        """Current power."""
-        return self.emeter_realtime.power
+    def current_power(self) -> float | None:
+        """Current power in watts."""
+        if power := self.energy.get("current_power"):
+            return power / 1_000
+        return None
 
     @property
     def energy(self):
@@ -72,22 +77,16 @@ def emeter_realtime(self):
         return EmeterStatus(
             {
                 "power_mw": self.energy.get("current_power"),
-                "total": self._convert_energy_data(
-                    self.energy.get("today_energy"), 1 / 1000
-                ),
+                "total": self.energy.get("today_energy") / 1_000,
             }
         )
 
     @property
     def emeter_this_month(self) -> float | None:
         """Get the emeter value for this month."""
-        return self._convert_energy_data(self.energy.get("month_energy"), 1 / 1000)
+        return self.energy.get("month_energy")
 
     @property
     def emeter_today(self) -> float | None:
         """Get the emeter value for today."""
-        return self._convert_energy_data(self.energy.get("today_energy"), 1 / 1000)
-
-    def _convert_energy_data(self, data, scale) -> float | None:
-        """Return adjusted emeter information."""
-        return data if not data else data * scale
+        return self.energy.get("today_energy")