feat：add missing parameter state_class (#101)

* feat: add state_class for sensor * feat: add optional state_class and unit_of_measurement for properties * feat:add support for state_class and unit process * style: fix pylint format * style: fix pylint format * sytle: change logic to suit conversation --------- Co-authored-by: LiShuzhen <SusanPhevos@gmail.com>
XiaoMi · Jan 3, 2025 · 5438698 · 5438698
1 parent da90e09
commit 5438698
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 5 deletions.
diff --git a/custom_components/xiaomi_home/miot/miot_device.py b/custom_components/xiaomi_home/miot/miot_device.py
@@ -515,9 +515,15 @@ def parse_miot_property_entity(
             prop.icon = self.icon_convert(prop.unit)
         device_class = SPEC_PROP_TRANS_MAP['properties'][prop_name][
             'device_class']
-        prop.platform = device_class
-
-        return {'platform': platform, 'device_class': device_class}
+        result = {'platform': platform, 'device_class': device_class}
+        # optional:
+        if 'optional' in SPEC_PROP_TRANS_MAP['properties'][prop_name]:
+            optional = SPEC_PROP_TRANS_MAP['properties'][prop_name]['optional']
+            if 'state_class' in optional:
+                result['state_class'] = optional['state_class']
+            if not prop.unit and 'unit_of_measurement' in optional:
+                result['unit_of_measurement'] = optional['unit_of_measurement']
+        return result
 
     def spec_transform(self) -> None:
         """Parse service, property, event, action from device spec."""
@@ -544,6 +550,13 @@ def spec_transform(self) -> None:
                 if prop_entity:
                     prop.platform = prop_entity['platform']
                     prop.device_class = prop_entity['device_class']
+                    if 'state_class' in prop_entity:
+                        prop.state_class = prop_entity['state_class']
+                    if 'unit_of_measurement' in prop_entity:
+                        prop.external_unit = self.unit_convert(
+                            prop_entity['unit_of_measurement'])
+                        prop.icon = self.icon_convert(
+                            prop_entity['unit_of_measurement'])
                 # general conversion
                 if not prop.platform:
                     if prop.writable:

diff --git a/custom_components/xiaomi_home/miot/miot_spec.py b/custom_components/xiaomi_home/miot/miot_spec.py
@@ -79,6 +79,7 @@ class MIoTSpecBase:
     # External params
     platform: str
     device_class: Any
+    state_class: Any
     icon: str
     external_unit: Any
 
@@ -96,6 +97,7 @@ def __init__(self, spec: dict) -> None:
 
         self.platform = None
         self.device_class = None
+        self.state_class = None
         self.icon = None
         self.external_unit = None
 

diff --git a/custom_components/xiaomi_home/miot/specs/specv2entity.py b/custom_components/xiaomi_home/miot/specs/specv2entity.py
@@ -46,8 +46,16 @@
 Conversion rules of MIoT-Spec-V2 instance to Home Assistant entity.
 """
 from homeassistant.components.sensor import SensorDeviceClass
+from homeassistant.components.sensor import SensorStateClass
 from homeassistant.components.event import EventDeviceClass
 
+from homeassistant.const import (
+    UnitOfEnergy,
+    UnitOfPower,
+    UnitOfElectricCurrent,
+    UnitOfElectricPotential,
+)
+
 # pylint: disable=pointless-string-statement
 """SPEC_DEVICE_TRANS_MAP
 {
@@ -325,7 +333,11 @@
     'properties': {
         '<property instance name>':{
             'device_class': str,
-            'entity': str
+            'entity': str,
+            'optional':{
+                'state_class': str,
+                'unit_of_measurement': str
+            }
         }
     }
 }
@@ -381,7 +393,11 @@
         },
         'voltage': {
             'device_class': SensorDeviceClass.VOLTAGE,
-            'entity': 'sensor'
+            'entity': 'sensor',
+            'optional': {
+                'state_class': SensorStateClass.MEASUREMENT,
+                'unit_of_measurement': UnitOfElectricPotential.VOLT
+            }
         },
         'illumination': {
             'device_class': SensorDeviceClass.ILLUMINANCE,
@@ -391,6 +407,38 @@
             'device_class': SensorDeviceClass.DURATION,
             'entity': 'sensor'
         },
+	    'electric-power': {
+            'device_class': SensorDeviceClass.POWER,
+            'entity': 'sensor',
+            'optional': {
+                'state_class': SensorStateClass.MEASUREMENT,
+                'unit_of_measurement': UnitOfPower.WATT
+            }
+	    },
+	    'electric-current': {
+            'device_class': SensorDeviceClass.CURRENT,
+            'entity': 'sensor',
+            'optional': {
+                'state_class': SensorStateClass.MEASUREMENT,
+                'unit_of_measurement': UnitOfElectricCurrent.AMPERE
+            }
+	    },
+	    'power-consumption': {
+            'device_class': SensorDeviceClass.ENERGY,
+            'entity': 'sensor',
+            'optional': {
+                'state_class': SensorStateClass.TOTAL_INCREASING,
+                'unit_of_measurement': UnitOfEnergy.KILO_WATT_HOUR
+            }
+	    },
+	    'total-battery': {
+            'device_class': SensorDeviceClass.ENERGY,
+            'entity': 'sensor',
+            'optional': {
+                'state_class': SensorStateClass.TOTAL_INCREASING,
+                'unit_of_measurement': UnitOfEnergy.KILO_WATT_HOUR
+            }
+	    },
         'has-someone-duration': 'no-one-determine-time',
         'no-one-duration': 'no-one-determine-time'
     }

diff --git a/custom_components/xiaomi_home/sensor.py b/custom_components/xiaomi_home/sensor.py
@@ -106,6 +106,9 @@ def __init__(self, miot_device: MIoTDevice, spec: MIoTSpecProperty) -> None:
         # Set icon
         if spec.icon:
             self._attr_icon = spec.icon
+        # Set state_class
+        if spec.state_class:
+            self._attr_state_class = spec.state_class
 
     @property
     def native_value(self) -> Any: