split entity module

custom_components/xiaomi_miot/__init__.py

@@ -47,7 +47,7 @@
     is_offline_exception,
     async_analytics_track_event,
 )
-from .core import HassEntry
+from .core import HassEntry, XEntity
 from .core.device import (
     Device,
     MiioDevice,
@@ -250,8 +250,7 @@ async def async_setup_entry(hass: hass_core.HomeAssistant, config_entry: config_
     else:
         entry = HassEntry.init(hass, config_entry)
         config = {**entry.get_config()}
-        device = entry.new_device(config)
-        await device.get_spec()
+        device = await entry.new_device(config)
         config[CONF_DEVICE] = device
         config[CONF_MODEL] = device.model
         config['miot_type'] = await device.get_urn()
@@ -294,10 +293,8 @@ async def async_setup_xiaomi_cloud(hass: hass_core.HomeAssistant, config_entry:
         cnt = len(config['devices_by_mac'])
         _LOGGER.debug('Setup xiaomi cloud for user: %s, %s devices', username, cnt)
     for mac, d in config['devices_by_mac'].items():
-        device = entry.new_device(d)
-        device.cloud = cloud
-        spec = await device.get_spec()
-        if not spec:
+        device = await entry.new_device(d, cloud)
+        if not device.spec:
             _LOGGER.warning('Xiaomi device: %s has no spec', device.name_model)
             continue
         conn = device.conn_mode
@@ -694,7 +691,9 @@ def get_device_class(self, enum):
 
     @property
     def model(self):
-        return self.device.info.model
+        if self.device:
+            return self.device.info.model
+        return self._model
 
     @property
     def name_model(self):

custom_components/xiaomi_miot/core/__init__.py

@@ -1,8 +1,10 @@
 from .device import Device, DeviceInfo
 from .hass_entry import HassEntry
+from .hass_entity import XEntity
 
 __all__ = [
     'Device',
     'DeviceInfo',
     'HassEntry',
+    'XEntity',
 ]

custom_components/xiaomi_miot/core/device.py

@@ -125,8 +125,6 @@ def __init__(self, info: DeviceInfo, entry: HassEntry):
         self.converters: list[BaseConv] = []
 
         self.local = MiotDevice.from_device(self)
-        if self.spec and not self.cloud_only:
-            self.miio2miot = Miio2MiotHelper.from_model(self.hass, self.model, self.spec)
 
     @cached_property
     def did(self):
@@ -260,6 +258,9 @@ def decode(self, data: dict | list) -> dict:
         return payload
 
     def decode_one(self, payload: dict, value: dict):
+        if not isinstance(value, dict):
+            _LOGGER.warning('%s: Device value is not dict: %s', self.name_model, value)
+            return
         if value.get('code', 0):
             return
         siid = value.get('siid')
@@ -291,6 +292,7 @@ async def update_miot_status(
     ) -> MiotResults:
         results = []
         self.miot_results = MiotResults(results)
+
         if use_local is None:
             use_local = self.local or self.miio2miot
             if self.cloud_only or use_cloud:

custom_components/xiaomi_miot/core/hass_entity.py

@@ -0,0 +1,63 @@
+import logging
+from typing import TYPE_CHECKING, Callable
+
+from homeassistant.helpers.entity import Entity, DeviceInfo
+
+from .miot_spec import MiotService, MiotProperty
+
+if TYPE_CHECKING:
+    from .device import Device
+    from .converters import BaseConv
+
+_LOGGER = logging.getLogger(__package__)
+
+class XEntity(Entity):
+    CLS: dict[str, Callable] = {}
+
+    log = _LOGGER
+    added = False
+    _attr_should_poll = False
+    _attr_has_entity_name = True
+
+    def __init__(self, device: 'Device', conv: 'BaseConv'):
+        self.device = device
+        self.hass = device.hass
+        self.conv = conv
+        attr = conv.attr
+        if isinstance(attr, MiotProperty):
+            self.attr = attr.full_name
+            self._attr_translation_key = attr.friendly_name
+            self.entity_id = attr.generate_entity_id(self, conv.domain)
+        elif isinstance(attr, MiotService):
+            self.attr = attr.name
+            self._attr_translation_key = attr.name
+            self.entity_id = attr.generate_entity_id(self, conv.domain)
+        else:
+            self.attr = attr
+            self._attr_translation_key = attr
+
+        self.on_init()
+
+    @property
+    def unique_mac(self):
+        return self.device.info.unique_id
+
+    def on_init(self):
+        """Run on class init."""
+
+    def on_device_update(self, data: dict):
+        pass
+
+    def get_state(self) -> dict:
+        """Run before entity remove if entity is subclass from RestoreEntity."""
+        return {}
+
+    def set_state(self, data: dict):
+        """Run on data from device."""
+        self._attr_state = data.get(self.attr)
+
+    async def async_added_to_hass(self) -> None:
+        self.device.add_listener(self.on_device_update)
+
+    async def async_will_remove_from_hass(self) -> None:
+        self.device.remove_listener(self.on_device_update)

custom_components/xiaomi_miot/core/hass_entry.py

@@ -3,6 +3,7 @@
 from homeassistant.const import CONF_USERNAME
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from .miio2miot import Miio2MiotHelper
 from .xiaomi_cloud import MiotCloud
 
 if TYPE_CHECKING:
@@ -38,10 +39,14 @@ def get_config(self, key=None, default=None):
             return dat.get(key, default)
         return dat
 
-    def new_device(self, device_info: dict):
+    async def new_device(self, device_info: dict, cloud: Optional[MiotCloud] = None):
         from .device import Device, DeviceInfo
         info = DeviceInfo(device_info)
         device = Device(info, self)
+        device.cloud = cloud
+        spec = await device.get_spec()
+        if spec and not device.cloud_only:
+            device.miio2miot = Miio2MiotHelper.from_model(self.hass, device.model, spec)
         self.devices[info.unique_id] = device
         return device
 

custom_components/xiaomi_miot/sensor.py

@@ -9,7 +9,7 @@
 from homeassistant.const import STATE_UNKNOWN
 from homeassistant.components.sensor import (
     DOMAIN as ENTITY_DOMAIN,
-    SensorEntity,
+    SensorEntity as BaseEntity,
     SensorDeviceClass,
     STATE_CLASSES,
 )
@@ -21,6 +21,7 @@
     CONF_MODEL,
     XIAOMI_CONFIG_SCHEMA as PLATFORM_SCHEMA,  # noqa: F401
     HassEntry,
+    XEntity,
     MiotEntity,
     BaseSubEntity,
     MiCoordinatorEntity,
@@ -140,7 +141,17 @@ def datetime_with_tzinfo(value):
     return value
 
 
-class MiotSensorEntity(MiotEntity, SensorEntity):
+class SensorEntity(XEntity, BaseEntity):
+    def get_state(self) -> dict:
+        return {self.attr: self._attr_native_value}
+
+    def set_state(self, data: dict):
+        self._attr_native_value = data.get(self.attr)
+
+XEntity.CLS[ENTITY_DOMAIN] = SensorEntity
+
+
+class MiotSensorEntity(MiotEntity, BaseEntity):
 
     def __init__(self, config, miot_service: MiotService):
         super().__init__(miot_service, config=config, logger=_LOGGER)
@@ -417,7 +428,7 @@ def turn_action(self, on):
         return ret
 
 
-class BaseSensorSubEntity(BaseSubEntity, SensorEntity):
+class BaseSensorSubEntity(BaseSubEntity, BaseEntity):
     def __init__(self, parent, attr, option=None, **kwargs):
         kwargs.setdefault('domain', ENTITY_DOMAIN)
         self._attr_state_class = None
@@ -507,7 +518,7 @@ def native_value(self):
         return val
 
 
-class MihomeMessageSensor(MiCoordinatorEntity, SensorEntity, RestoreEntity):
+class MihomeMessageSensor(MiCoordinatorEntity, BaseEntity, RestoreEntity):
     _filter_homes = None
     _exclude_types = None
     _has_none_message = False
@@ -640,7 +651,7 @@ async def fetch_latest_message(self):
         return msg
 
 
-class MihomeSceneHistorySensor(MiCoordinatorEntity, SensorEntity, RestoreEntity):
+class MihomeSceneHistorySensor(MiCoordinatorEntity, BaseEntity, RestoreEntity):
     MESSAGE_TIMEOUT = 60
     UPDATE_INTERVAL = 15