From 0ad10eba7b60df5e1df8067deac746222c7520e5 Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Thu, 20 Jun 2024 02:01:29 +0200
Subject: [PATCH] Add unit_getter for feature

---
 kasa/feature.py                         | 9 ++++++++-
 kasa/smart/modules/temperaturesensor.py | 4 ++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/kasa/feature.py b/kasa/feature.py
index d0c83a3dc..38c6fca99 100644
--- a/kasa/feature.py
+++ b/kasa/feature.py
@@ -99,7 +99,7 @@ class Type(Enum):
         Choice = auto()
         Unknown = -1
 
-    # TODO: unsure if this is a great idea..
+    # Aliases for easy access
     Sensor = Type.Sensor
     BinarySensor = Type.BinarySensor
     Switch = Type.Switch
@@ -139,6 +139,9 @@ class Category(Enum):
     icon: str | None = None
     #: Unit, if applicable
     unit: str | None = None
+    #: Attribute containing the name of the unit getter property.
+    #: If set, this property will be used to set *unit*.
+    unit_getter: str | None = None
     #: Category hint for downstreams
     category: Feature.Category = Category.Unset
     #: Type of the feature
@@ -177,6 +180,10 @@ def __post_init__(self):
         if self.choices_getter is not None:
             self.choices = getattr(container, self.choices_getter)
 
+        # Populate unit, if unit_getter is given
+        if self.unit_getter is not None:
+            self.unit = getattr(container, self.unit_getter)
+
         # Set the category, if unset
         if self.category is Feature.Category.Unset:
             if self.attribute_setter:
diff --git a/kasa/smart/modules/temperaturesensor.py b/kasa/smart/modules/temperaturesensor.py
index 4880fc301..d58ffd235 100644
--- a/kasa/smart/modules/temperaturesensor.py
+++ b/kasa/smart/modules/temperaturesensor.py
@@ -28,6 +28,7 @@ def __init__(self, device: SmartDevice, module: str):
                 attribute_getter="temperature",
                 icon="mdi:thermometer",
                 category=Feature.Category.Primary,
+                unit_getter="temperature_unit",
             )
         )
         if "current_temp_exception" in device.sys_info:
@@ -55,7 +56,6 @@ def __init__(self, device: SmartDevice, module: str):
                 choices=["celsius", "fahrenheit"],
             )
         )
-        # TODO: use temperature_unit for feature creation
 
     @property
     def temperature(self):
@@ -68,7 +68,7 @@ def temperature_warning(self) -> bool:
         return self._device.sys_info.get("current_temp_exception", 0) != 0
 
     @property
-    def temperature_unit(self):
+    def temperature_unit(self) -> Literal["celsius", "fahrenheit"]:
         """Return current temperature unit."""
         return self._device.sys_info["temp_unit"]