From f69ea6017d67e4937ee0570a5241be4fafb2690b Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Tue, 14 Jun 2022 16:00:07 +0200
Subject: [PATCH] Add device class support to Tuya number entities (#73483)

---
 homeassistant/components/tuya/number.py | 49 +++++++++++++++++++++++--
 1 file changed, 45 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/tuya/number.py b/homeassistant/components/tuya/number.py
index 35efd78871f..a342fe58bd2 100644
--- a/homeassistant/components/tuya/number.py
+++ b/homeassistant/components/tuya/number.py
@@ -3,7 +3,11 @@ from __future__ import annotations
 
 from tuya_iot import TuyaDevice, TuyaDeviceManager
 
-from homeassistant.components.number import NumberEntity, NumberEntityDescription
+from homeassistant.components.number import (
+    NumberDeviceClass,
+    NumberEntity,
+    NumberEntityDescription,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -12,7 +16,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import HomeAssistantTuyaData
 from .base import IntegerTypeData, TuyaEntity
-from .const import DOMAIN, TUYA_DISCOVERY_NEW, DPCode, DPType
+from .const import DEVICE_CLASS_UNITS, DOMAIN, TUYA_DISCOVERY_NEW, DPCode, DPType
 
 # All descriptions can be found here. Mostly the Integer data types in the
 # default instructions set of each category end up being a number.
@@ -33,24 +37,28 @@ NUMBERS: dict[str, tuple[NumberEntityDescription, ...]] = {
         NumberEntityDescription(
             key=DPCode.TEMP_SET,
             name="Temperature",
+            device_class=NumberDeviceClass.TEMPERATURE,
             icon="mdi:thermometer",
             entity_category=EntityCategory.CONFIG,
         ),
         NumberEntityDescription(
             key=DPCode.TEMP_SET_F,
             name="Temperature",
+            device_class=NumberDeviceClass.TEMPERATURE,
             icon="mdi:thermometer",
             entity_category=EntityCategory.CONFIG,
         ),
         NumberEntityDescription(
             key=DPCode.TEMP_BOILING_C,
             name="Temperature After Boiling",
+            device_class=NumberDeviceClass.TEMPERATURE,
             icon="mdi:thermometer",
             entity_category=EntityCategory.CONFIG,
         ),
         NumberEntityDescription(
             key=DPCode.TEMP_BOILING_F,
             name="Temperature After Boiling",
+            device_class=NumberDeviceClass.TEMPERATURE,
             icon="mdi:thermometer",
             entity_category=EntityCategory.CONFIG,
         ),
@@ -108,6 +116,7 @@ NUMBERS: dict[str, tuple[NumberEntityDescription, ...]] = {
         NumberEntityDescription(
             key=DPCode.TEMP_SET,
             name="Temperature",
+            device_class=NumberDeviceClass.TEMPERATURE,
             icon="mdi:thermometer",
             entity_category=EntityCategory.CONFIG,
         ),
@@ -256,6 +265,7 @@ NUMBERS: dict[str, tuple[NumberEntityDescription, ...]] = {
         NumberEntityDescription(
             key=DPCode.TEMP,
             name="Temperature",
+            device_class=NumberDeviceClass.TEMPERATURE,
             icon="mdi:thermometer-lines",
         ),
     ),
@@ -265,11 +275,13 @@ NUMBERS: dict[str, tuple[NumberEntityDescription, ...]] = {
         NumberEntityDescription(
             key=DPCode.TEMP_SET,
             name="Temperature",
+            device_class=NumberDeviceClass.TEMPERATURE,
             icon="mdi:thermometer-lines",
         ),
         NumberEntityDescription(
             key=DPCode.TEMP_SET_F,
             name="Temperature",
+            device_class=NumberDeviceClass.TEMPERATURE,
             icon="mdi:thermometer-lines",
         ),
     ),
@@ -329,8 +341,37 @@ class TuyaNumberEntity(TuyaEntity, NumberEntity):
             self._attr_max_value = self._number.max_scaled
             self._attr_min_value = self._number.min_scaled
             self._attr_step = self._number.step_scaled
-            if description.unit_of_measurement is None:
-                self._attr_unit_of_measurement = self._number.unit
+
+        # Logic to ensure the set device class and API received Unit Of Measurement
+        # match Home Assistants requirements.
+        if (
+            self.device_class is not None
+            and not self.device_class.startswith(DOMAIN)
+            and description.native_unit_of_measurement is None
+        ):
+
+            # We cannot have a device class, if the UOM isn't set or the
+            # device class cannot be found in the validation mapping.
+            if (
+                self.native_unit_of_measurement is None
+                or self.device_class not in DEVICE_CLASS_UNITS
+            ):
+                self._attr_device_class = None
+                return
+
+            uoms = DEVICE_CLASS_UNITS[self.device_class]
+            self._uom = uoms.get(self.native_unit_of_measurement) or uoms.get(
+                self.native_unit_of_measurement.lower()
+            )
+
+            # Unknown unit of measurement, device class should not be used.
+            if self._uom is None:
+                self._attr_device_class = None
+                return
+
+            # If we still have a device class, we should not use an icon
+            if self.device_class:
+                self._attr_icon = None
 
     @property
     def value(self) -> float | None: