From 3b212b9109817929da3971ebd0d57849a83d7878 Mon Sep 17 00:00:00 2001
From: Maciej Bieniek <bieniu@users.noreply.github.com>
Date: Wed, 4 Aug 2021 12:03:18 +0200
Subject: [PATCH] Use `NumberEntityDescription` for Xiaomi Miio (#53911)

---
 .../components/xiaomi_miio/number.py          | 69 +++++++++----------
 1 file changed, 32 insertions(+), 37 deletions(-)

diff --git a/homeassistant/components/xiaomi_miio/number.py b/homeassistant/components/xiaomi_miio/number.py
index 6855faa6391..83c336d101e 100644
--- a/homeassistant/components/xiaomi_miio/number.py
+++ b/homeassistant/components/xiaomi_miio/number.py
@@ -1,8 +1,10 @@
 """Motor speed support for Xiaomi Mi Air Humidifier."""
+from __future__ import annotations
+
 from dataclasses import dataclass
 from enum import Enum
 
-from homeassistant.components.number import NumberEntity
+from homeassistant.components.number import NumberEntity, NumberEntityDescription
 from homeassistant.core import callback
 
 from .const import (
@@ -21,28 +23,23 @@ ATTR_MOTOR_SPEED = "motor_speed"
 
 
 @dataclass
-class NumberType:
-    """Class that holds device specific info for a xiaomi aqara or humidifier number controller types."""
+class XiaomiMiioNumberDescription(NumberEntityDescription):
+    """A class that describes number entities."""
 
-    name: str = None
-    short_name: str = None
-    unit_of_measurement: str = None
-    icon: str = None
-    device_class: str = None
-    min: float = None
-    max: float = None
-    step: float = None
+    min_value: float | None = None
+    max_value: float | None = None
+    step: float | None = None
     available_with_device_off: bool = True
 
 
 NUMBER_TYPES = {
-    FEATURE_SET_MOTOR_SPEED: NumberType(
+    FEATURE_SET_MOTOR_SPEED: XiaomiMiioNumberDescription(
+        key=ATTR_MOTOR_SPEED,
         name="Motor Speed",
         icon="mdi:fast-forward-outline",
-        short_name=ATTR_MOTOR_SPEED,
         unit_of_measurement="rpm",
-        min=200,
-        max=2000,
+        min_value=200,
+        max_value=2000,
         step=10,
         available_with_device_off=False,
     ),
@@ -56,22 +53,21 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         return
     model = config_entry.data[CONF_MODEL]
     device = hass.data[DOMAIN][config_entry.entry_id][KEY_DEVICE]
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][KEY_COORDINATOR]
 
-    if model not in [MODEL_AIRHUMIDIFIER_CA4]:
+    if model != MODEL_AIRHUMIDIFIER_CA4:
         return
 
-    for number in NUMBER_TYPES.values():
-        entities.append(
-            XiaomiAirHumidifierNumber(
-                f"{config_entry.title} {number.name}",
-                device,
-                config_entry,
-                f"{number.short_name}_{config_entry.unique_id}",
-                number,
-                coordinator,
-            )
+    description = NUMBER_TYPES[FEATURE_SET_MOTOR_SPEED]
+    entities.append(
+        XiaomiAirHumidifierNumber(
+            f"{config_entry.title} {description.name}",
+            device,
+            config_entry,
+            f"{description.key}_{config_entry.unique_id}",
+            hass.data[DOMAIN][config_entry.entry_id][KEY_COORDINATOR],
+            description,
         )
+    )
 
     async_add_entities(entities)
 
@@ -79,18 +75,17 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 class XiaomiAirHumidifierNumber(XiaomiCoordinatedMiioEntity, NumberEntity):
     """Representation of a generic Xiaomi attribute selector."""
 
-    def __init__(self, name, device, entry, unique_id, number, coordinator):
+    def __init__(self, name, device, entry, unique_id, coordinator, description):
         """Initialize the generic Xiaomi attribute selector."""
         super().__init__(name, device, entry, unique_id, coordinator)
-        self._attr_icon = number.icon
-        self._attr_unit_of_measurement = number.unit_of_measurement
-        self._attr_min_value = number.min
-        self._attr_max_value = number.max
-        self._attr_step = number.step
-        self._controller = number
+
+        self._attr_min_value = description.min_value
+        self._attr_max_value = description.max_value
+        self._attr_step = description.step
         self._attr_value = self._extract_value_from_attribute(
-            self.coordinator.data, self._controller.short_name
+            coordinator.data, description.key
         )
+        self.entity_description = description
 
     @property
     def available(self):
@@ -98,7 +93,7 @@ class XiaomiAirHumidifierNumber(XiaomiCoordinatedMiioEntity, NumberEntity):
         if (
             super().available
             and not self.coordinator.data.is_on
-            and not self._controller.available_with_device_off
+            and not self.entity_description.available_with_device_off
         ):
             return False
         return super().available
@@ -131,7 +126,7 @@ class XiaomiAirHumidifierNumber(XiaomiCoordinatedMiioEntity, NumberEntity):
         """Fetch state from the device."""
         # On state change the device doesn't provide the new state immediately.
         self._attr_value = self._extract_value_from_attribute(
-            self.coordinator.data, self._controller.short_name
+            self.coordinator.data, self.entity_description.key
         )
         self.async_write_ha_state()