From 0fc92928a4544a3f4d3021f7a0923e7b950c4e84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vladim=C3=ADr=20Z=C3=A1hradn=C3=ADk?=
 <vladimir@zahradnik.io>
Date: Sat, 21 Dec 2019 10:54:48 +0100
Subject: [PATCH] Add device class attribute to modbus sensors (#30030)

---
 .../components/modbus/binary_sensor.py        | 24 +++++++++++++++----
 homeassistant/components/modbus/sensor.py     | 14 +++++++++--
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/modbus/binary_sensor.py b/homeassistant/components/modbus/binary_sensor.py
index faf5160d7e5..9a431d24b0c 100644
--- a/homeassistant/components/modbus/binary_sensor.py
+++ b/homeassistant/components/modbus/binary_sensor.py
@@ -1,10 +1,15 @@
 """Support for Modbus Coil sensors."""
 import logging
+from typing import Optional
 
 import voluptuous as vol
 
-from homeassistant.components.binary_sensor import PLATFORM_SCHEMA, BinarySensorDevice
-from homeassistant.const import CONF_NAME, CONF_SLAVE
+from homeassistant.components.binary_sensor import (
+    DEVICE_CLASSES_SCHEMA,
+    PLATFORM_SCHEMA,
+    BinarySensorDevice,
+)
+from homeassistant.const import CONF_DEVICE_CLASS, CONF_NAME, CONF_SLAVE
 from homeassistant.helpers import config_validation as cv
 
 from . import CONF_HUB, DEFAULT_HUB, DOMAIN as MODBUS_DOMAIN
@@ -20,6 +25,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
             {
                 vol.Required(CONF_COIL): cv.positive_int,
                 vol.Required(CONF_NAME): cv.string,
+                vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
                 vol.Optional(CONF_HUB, default=DEFAULT_HUB): cv.string,
                 vol.Optional(CONF_SLAVE): cv.positive_int,
             }
@@ -35,7 +41,11 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
         hub = hass.data[MODBUS_DOMAIN][coil.get(CONF_HUB)]
         sensors.append(
             ModbusCoilSensor(
-                hub, coil.get(CONF_NAME), coil.get(CONF_SLAVE), coil.get(CONF_COIL)
+                hub,
+                coil.get(CONF_NAME),
+                coil.get(CONF_SLAVE),
+                coil.get(CONF_COIL),
+                coil.get(CONF_DEVICE_CLASS),
             )
         )
 
@@ -45,12 +55,13 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
 class ModbusCoilSensor(BinarySensorDevice):
     """Modbus coil sensor."""
 
-    def __init__(self, hub, name, slave, coil):
+    def __init__(self, hub, name, slave, coil, device_class):
         """Initialize the Modbus coil sensor."""
         self._hub = hub
         self._name = name
         self._slave = int(slave) if slave else None
         self._coil = int(coil)
+        self._device_class = device_class
         self._value = None
 
     @property
@@ -63,6 +74,11 @@ class ModbusCoilSensor(BinarySensorDevice):
         """Return the state of the sensor."""
         return self._value
 
+    @property
+    def device_class(self) -> Optional[str]:
+        """Return the device class of the sensor."""
+        return self._device_class
+
     def update(self):
         """Update the state of the sensor."""
         result = self._hub.read_coils(self._slave, self._coil, 1)
diff --git a/homeassistant/components/modbus/sensor.py b/homeassistant/components/modbus/sensor.py
index 86f6445b8d6..5b04a898ab9 100644
--- a/homeassistant/components/modbus/sensor.py
+++ b/homeassistant/components/modbus/sensor.py
@@ -1,12 +1,13 @@
 """Support for Modbus Register sensors."""
 import logging
 import struct
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA
+from homeassistant.components.sensor import DEVICE_CLASSES_SCHEMA, PLATFORM_SCHEMA
 from homeassistant.const import (
+    CONF_DEVICE_CLASS,
     CONF_NAME,
     CONF_OFFSET,
     CONF_SLAVE,
@@ -67,6 +68,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
                 vol.Optional(CONF_DATA_TYPE, default=DATA_TYPE_INT): vol.In(
                     [DATA_TYPE_INT, DATA_TYPE_UINT, DATA_TYPE_FLOAT, DATA_TYPE_CUSTOM]
                 ),
+                vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
                 vol.Optional(CONF_HUB, default=DEFAULT_HUB): cv.string,
                 vol.Optional(CONF_OFFSET, default=0): number,
                 vol.Optional(CONF_PRECISION, default=0): cv.positive_int,
@@ -139,6 +141,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
                 register.get(CONF_OFFSET),
                 structure,
                 register.get(CONF_PRECISION),
+                register.get(CONF_DEVICE_CLASS),
             )
         )
 
@@ -164,6 +167,7 @@ class ModbusRegisterSensor(RestoreEntity):
         offset,
         structure,
         precision,
+        device_class,
     ):
         """Initialize the modbus register sensor."""
         self._hub = hub
@@ -178,6 +182,7 @@ class ModbusRegisterSensor(RestoreEntity):
         self._offset = offset
         self._precision = precision
         self._structure = structure
+        self._device_class = device_class
         self._value = None
 
     async def async_added_to_hass(self):
@@ -202,6 +207,11 @@ class ModbusRegisterSensor(RestoreEntity):
         """Return the unit of measurement."""
         return self._unit_of_measurement
 
+    @property
+    def device_class(self) -> Optional[str]:
+        """Return the device class of the sensor."""
+        return self._device_class
+
     def update(self):
         """Update the state of the sensor."""
         if self._register_type == REGISTER_TYPE_INPUT: