Add support for feature units (#843)

rytilahti · sdb9696 · web-flow · commit 890900daf330 · 2024-04-22T10:25:30.000+01:00
Co-authored-by: Steven B &lt;51370195+sdb9696@users.noreply.github.com&gt;
diff --git a/kasa/cli.py b/kasa/cli.py
@@ -598,9 +598,10 @@ async def state(ctx, dev: Device):
         echo("\t[bold]== Children ==[/bold]")
         for child in dev.children:
             echo(f"\t* {child.alias} ({child.model}, {child.device_type})")
-            for feat in child.features.values():
+            for id_, feat in child.features.items():
                 try:
-                    echo(f"\t\t{feat.name}: {feat.value}")
+                    unit = f" {feat.unit}" if feat.unit else ""
+                    echo(f"\t\t{feat.name} ({id_}): {feat.value}{unit}")
                 except Exception as ex:
                     echo(f"\t\t{feat.name}: got exception (%s)" % ex)
         echo()
@@ -614,12 +615,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():
@@ -1177,7 +1174,8 @@ async def feature(dev: Device, child: str, name: str, value):
         def _print_features(dev):
             for name, feat in dev.features.items():
                 try:
-                    echo(f"\t{feat.name} ({name}): {feat.value}")
+                    unit = f" {feat.unit}" if feat.unit else ""
+                    echo(f"\t{feat.name} ({name}): {feat.value}{unit}")
                 except Exception as ex:
                     echo(f"\t{feat.name} ({name}): [red]{ex}[/red]")
 
@@ -1198,7 +1196,8 @@ def _print_features(dev):
     feat = dev.features[name]
 
     if value is None:
-        echo(f"{feat.name} ({name}): {feat.value}")
+        unit = f" {feat.unit}" if feat.unit else ""
+        echo(f"{feat.name} ({name}): {feat.value}{unit}")
         return feat.value
 
     echo(f"Setting {name} to {value}")
diff --git 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
@@ -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")
diff --git a/kasa/tests/test_feature.py b/kasa/tests/test_feature.py
@@ -17,6 +17,7 @@ class DummyDevice:
         container=None,
         icon="mdi:dummy",
         type=FeatureType.BinarySensor,
+        unit="dummyunit",
     )
     return feat
 
@@ -30,6 +31,7 @@ def test_feature_api(dummy_feature: Feature):
     assert dummy_feature.container is None
     assert dummy_feature.icon == "mdi:dummy"
     assert dummy_feature.type == FeatureType.BinarySensor
+    assert dummy_feature.unit == "dummyunit"
 
 
 def test_feature_value(dummy_feature: Feature):
