Convert entity class from str to enum #1316

AlexxIT · Apr 4, 2024 · 534128c · 534128c
1 parent 23e0055
commit 534128c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/custom_components/xiaomi_gateway3/hass/entity_description.py b/custom_components/xiaomi_gateway3/hass/entity_description.py
@@ -225,8 +225,11 @@ def setup_entity_description(entity: Entity, conv: BaseConv) -> bool:
     for k, v in desc.items():
         if k == "lazy" or v is None:
             continue
-        if k == "category":
+        if k == "category" and type(v) is str:
             v = EntityCategory(v)
+        elif k == "class" and type(v) is str:
+            if domain_class := DOMAIN_CLASSES.get(conv.domain):
+                v = domain_class(v)
         setattr(entity, ENTITY_KEYS.get(k) or k, v)
 
     # sensor with unit_of_measurement and without state_class will be MEASUREMENT