Device class attribute for sensor component (#1525)

* Add constants for sensor device_class

* Add device_class attribute to sensor component

* Add device_class attribute to sensor class

* Add device_class to mhz19 temperature sensor

* Add device_class to sensor in api component

* Add test for device_class of sensor

* Rename DEVICE_CLASS_NONE to DEVICE_CLASS_EMPTY for consistency

* Make optional attributes of sensor component truly optional

esphome/components/api/api.proto

@@ -419,6 +419,7 @@ message ListEntitiesSensorResponse {
   string unit_of_measurement = 6;
   int32 accuracy_decimals = 7;
   bool force_update = 8;
+  string device_class = 9;
 }
 message SensorStateResponse {
   option (id) = 25;

esphome/components/api/api_connection.cpp

@@ -382,6 +382,7 @@ bool APIConnection::send_sensor_info(sensor::Sensor *sensor) {
   msg.unit_of_measurement = sensor->get_unit_of_measurement();
   msg.accuracy_decimals = sensor->get_accuracy_decimals();
   msg.force_update = sensor->get_force_update();
+  msg.device_class = sensor->get_device_class();
   return this->send_list_entities_sensor_response(msg);
 }
 #endif

esphome/components/api/api_pb2.cpp

@@ -1494,6 +1494,10 @@ bool ListEntitiesSensorResponse::decode_length(uint32_t field_id, ProtoLengthDel
       this->unit_of_measurement = value.as_string();
       return true;
     }
+    case 9: {
+      this->device_class = value.as_string();
+      return true;
+    }
     default:
       return false;
   }
@@ -1517,6 +1521,7 @@ void ListEntitiesSensorResponse::encode(ProtoWriteBuffer buffer) const {
   buffer.encode_string(6, this->unit_of_measurement);
   buffer.encode_int32(7, this->accuracy_decimals);
   buffer.encode_bool(8, this->force_update);
+  buffer.encode_string(9, this->device_class);
 }
 void ListEntitiesSensorResponse::dump_to(std::string &out) const {
   char buffer[64];
@@ -1554,6 +1559,10 @@ void ListEntitiesSensorResponse::dump_to(std::string &out) const {
   out.append("  force_update: ");
   out.append(YESNO(this->force_update));
   out.append("\n");
+
+  out.append("  device_class: ");
+  out.append("'").append(this->device_class).append("'");
+  out.append("\n");
   out.append("}");
 }
 bool SensorStateResponse::decode_varint(uint32_t field_id, ProtoVarInt value) {

esphome/components/api/api_pb2.h

@@ -401,6 +401,7 @@ class ListEntitiesSensorResponse : public ProtoMessage {
   std::string unique_id{};            // NOLINT
   std::string icon{};                 // NOLINT
   std::string unit_of_measurement{};  // NOLINT
+  std::string device_class{};         // NOLINT
   int32_t accuracy_decimals{0};       // NOLINT
   bool force_update{false};           // NOLINT
   void encode(ProtoWriteBuffer buffer) const override;

esphome/components/mhz19/sensor.py

@@ -4,7 +4,7 @@ from esphome import automation
 from esphome.automation import maybe_simple_id
 from esphome.components import sensor, uart
 from esphome.const import CONF_CO2, CONF_ID, CONF_TEMPERATURE, ICON_MOLECULE_CO2, \
-    UNIT_PARTS_PER_MILLION, UNIT_CELSIUS, ICON_THERMOMETER
+    UNIT_PARTS_PER_MILLION, UNIT_CELSIUS, ICON_EMPTY, DEVICE_CLASS_TEMPERATURE
 
 DEPENDENCIES = ['uart']
 
@@ -19,7 +19,8 @@ MHZ19ABCDisableAction = mhz19_ns.class_('MHZ19ABCDisableAction', automation.Acti
 CONFIG_SCHEMA = cv.Schema({
     cv.GenerateID(): cv.declare_id(MHZ19Component),
     cv.Required(CONF_CO2): sensor.sensor_schema(UNIT_PARTS_PER_MILLION, ICON_MOLECULE_CO2, 0),
-    cv.Optional(CONF_TEMPERATURE): sensor.sensor_schema(UNIT_CELSIUS, ICON_THERMOMETER, 0),
+    cv.Optional(CONF_TEMPERATURE): sensor.sensor_schema(
+        UNIT_CELSIUS, ICON_EMPTY, 0, DEVICE_CLASS_TEMPERATURE),
     cv.Optional(CONF_AUTOMATIC_BASELINE_CALIBRATION): cv.boolean,
 }).extend(cv.polling_component_schema('60s')).extend(uart.UART_DEVICE_SCHEMA)
 

esphome/components/sensor/__init__.py

@@ -4,15 +4,25 @@ import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome import automation
 from esphome.components import mqtt
-from esphome.const import CONF_ABOVE, CONF_ACCURACY_DECIMALS, CONF_ALPHA, CONF_BELOW, \
-    CONF_EXPIRE_AFTER, CONF_FILTERS, CONF_FROM, CONF_ICON, CONF_ID, CONF_INTERNAL, \
+from esphome.const import CONF_DEVICE_CLASS, CONF_ABOVE, CONF_ACCURACY_DECIMALS, CONF_ALPHA, \
+    CONF_BELOW, CONF_EXPIRE_AFTER, CONF_FILTERS, CONF_FROM, CONF_ICON, CONF_ID, CONF_INTERNAL, \
     CONF_ON_RAW_VALUE, CONF_ON_VALUE, CONF_ON_VALUE_RANGE, CONF_SEND_EVERY, CONF_SEND_FIRST_AT, \
     CONF_TO, CONF_TRIGGER_ID, CONF_UNIT_OF_MEASUREMENT, CONF_WINDOW_SIZE, CONF_NAME, CONF_MQTT_ID, \
-    CONF_FORCE_UPDATE
+    CONF_FORCE_UPDATE, UNIT_EMPTY, ICON_EMPTY, DEVICE_CLASS_EMPTY, DEVICE_CLASS_BATTERY, \
+    DEVICE_CLASS_CURRENT, DEVICE_CLASS_ENERGY, DEVICE_CLASS_HUMIDITY, DEVICE_CLASS_ILLUMINANCE, \
+    DEVICE_CLASS_SIGNAL_STRENGTH, DEVICE_CLASS_TEMPERATURE, DEVICE_CLASS_POWER, \
+    DEVICE_CLASS_POWER_FACTOR, DEVICE_CLASS_PRESSURE, DEVICE_CLASS_TIMESTAMP, DEVICE_CLASS_VOLTAGE
 from esphome.core import CORE, coroutine, coroutine_with_priority
 from esphome.util import Registry
 
 CODEOWNERS = ['@esphome/core']
+DEVICE_CLASSES = [
+    DEVICE_CLASS_EMPTY, DEVICE_CLASS_BATTERY, DEVICE_CLASS_CURRENT, DEVICE_CLASS_ENERGY,
+    DEVICE_CLASS_HUMIDITY, DEVICE_CLASS_ILLUMINANCE, DEVICE_CLASS_SIGNAL_STRENGTH,
+    DEVICE_CLASS_TEMPERATURE, DEVICE_CLASS_POWER, DEVICE_CLASS_POWER_FACTOR, DEVICE_CLASS_PRESSURE,
+    DEVICE_CLASS_TIMESTAMP, DEVICE_CLASS_VOLTAGE
+]
+
 IS_PLATFORM_COMPONENT = True
 
 
@@ -80,6 +90,7 @@ SensorInRangeCondition = sensor_ns.class_('SensorInRangeCondition', Filter)
 unit_of_measurement = cv.string_strict
 accuracy_decimals = cv.int_
 icon = cv.icon
+device_class = cv.one_of(*DEVICE_CLASSES, lower=True, space='_')
 
 SENSOR_SCHEMA = cv.MQTT_COMPONENT_SCHEMA.extend({
     cv.OnlyWith(CONF_MQTT_ID, 'mqtt'): cv.declare_id(mqtt.MQTTSensorComponent),
@@ -87,6 +98,7 @@ SENSOR_SCHEMA = cv.MQTT_COMPONENT_SCHEMA.extend({
     cv.Optional(CONF_UNIT_OF_MEASUREMENT): unit_of_measurement,
     cv.Optional(CONF_ICON): icon,
     cv.Optional(CONF_ACCURACY_DECIMALS): accuracy_decimals,
+    cv.Optional(CONF_DEVICE_CLASS): device_class,
     cv.Optional(CONF_FORCE_UPDATE, default=False): cv.boolean,
     cv.Optional(CONF_EXPIRE_AFTER): cv.All(cv.requires_component('mqtt'),
                                            cv.Any(None, cv.positive_time_period_milliseconds)),
@@ -105,13 +117,25 @@ SENSOR_SCHEMA = cv.MQTT_COMPONENT_SCHEMA.extend({
 })
 
 
-def sensor_schema(unit_of_measurement_, icon_, accuracy_decimals_):
-    # type: (str, str, int) -> cv.Schema
-    return SENSOR_SCHEMA.extend({
-        cv.Optional(CONF_UNIT_OF_MEASUREMENT, default=unit_of_measurement_): unit_of_measurement,
-        cv.Optional(CONF_ICON, default=icon_): icon,
+def sensor_schema(unit_of_measurement_=UNIT_EMPTY, icon_=ICON_EMPTY, accuracy_decimals_=0,
+                  device_class_=DEVICE_CLASS_EMPTY):
+    # type: (str, str, int, str) -> cv.Schema
+    schema = SENSOR_SCHEMA
+    if unit_of_measurement_ != UNIT_EMPTY:
+        schema = schema.extend({
+            cv.Optional(CONF_UNIT_OF_MEASUREMENT, default=unit_of_measurement_): unit_of_measurement
+        })
+    if icon_ != ICON_EMPTY:
+        schema = schema.extend({cv.Optional(CONF_ICON, default=icon_): icon})
+    if accuracy_decimals_ != 0:
+        schema = schema.extend({
             cv.Optional(CONF_ACCURACY_DECIMALS, default=accuracy_decimals_): accuracy_decimals,
         })
+    if device_class_ != DEVICE_CLASS_EMPTY:
+        schema = schema.extend({
+            cv.Optional(CONF_DEVICE_CLASS, default=device_class_): device_class
+        })
+    return schema
 
 
 @FILTER_REGISTRY.register('offset', OffsetFilter, cv.float_)
@@ -253,6 +277,8 @@ def setup_sensor_core_(var, config):
     cg.add(var.set_name(config[CONF_NAME]))
     if CONF_INTERNAL in config:
         cg.add(var.set_internal(config[CONF_INTERNAL]))
+    if CONF_DEVICE_CLASS in config:
+        cg.add(var.set_device_class(config[CONF_DEVICE_CLASS]))
     if CONF_UNIT_OF_MEASUREMENT in config:
         cg.add(var.set_unit_of_measurement(config[CONF_UNIT_OF_MEASUREMENT]))
     if CONF_ICON in config:

esphome/components/sensor/sensor.cpp

@@ -40,6 +40,13 @@ std::string Sensor::get_icon() {
     return *this->icon_;
   return this->icon();
 }
+void Sensor::set_device_class(const std::string &device_class) { this->device_class_ = device_class; }
+std::string Sensor::get_device_class() {
+  if (this->device_class_.has_value())
+    return *this->device_class_;
+  return this->device_class();
+}
+std::string Sensor::device_class() { return ""; }
 std::string Sensor::get_unit_of_measurement() {
   if (this->unit_of_measurement_.has_value())
     return *this->unit_of_measurement_;

esphome/components/sensor/sensor.h

@@ -10,6 +10,9 @@ namespace sensor {
 #define LOG_SENSOR(prefix, type, obj) \
   if (obj != nullptr) { \
     ESP_LOGCONFIG(TAG, "%s%s '%s'", prefix, type, obj->get_name().c_str()); \
+    if (!obj->get_device_class().empty()) { \
+      ESP_LOGCONFIG(TAG, "%s  Device Class: '%s'", prefix, obj->get_device_class().c_str()); \
+    } \
     ESP_LOGCONFIG(TAG, "%s  Unit of Measurement: '%s'", prefix, obj->get_unit_of_measurement().c_str()); \
     ESP_LOGCONFIG(TAG, "%s  Accuracy Decimals: %d", prefix, obj->get_accuracy_decimals()); \
     if (!obj->get_icon().empty()) { \
@@ -122,6 +125,12 @@ class Sensor : public Nameable {
    */
   float state;
 
+  /// Manually set the Home Assistant device class (see sensor::device_class)
+  void set_device_class(const std::string &device_class);
+
+  /// Get the device class for this sensor, using the manual override if specified.
+  std::string get_device_class();
+
   /** This member variable stores the current raw state of the sensor. Unlike .state,
    * this will be updated immediately when publish_state is called.
    */
@@ -130,6 +139,14 @@ class Sensor : public Nameable {
   /// Return whether this sensor has gotten a full state (that passed through all filters) yet.
   bool has_state() const;
 
+  /** Override this to set the Home Assistant device class for this sensor.
+   *
+   * Return "" to disable this feature.
+   *
+   * @return The device class of this sensor, for example "temperature".
+   */
+  virtual std::string device_class();
+
   /** A unique ID for this sensor, empty for no unique id. See unique ID requirements:
    * https://developers.home-assistant.io/docs/en/entity_registry_index.html#unique-id-requirements
    *
@@ -174,6 +191,8 @@ class Sensor : public Nameable {
   /// Return the accuracy in decimals for this sensor.
   virtual int8_t accuracy_decimals();  // NOLINT
 
+  optional<std::string> device_class_{};  ///< Stores the override of the device class
+
   uint32_t hash_base() override;
 
   CallbackManager<void(float)> raw_callback_;  ///< Storage for raw state callbacks.

esphome/const.py

@@ -663,3 +663,17 @@ UNIT_WATT_HOURS = 'Wh'
 
 DEVICE_CLASS_CONNECTIVITY = 'connectivity'
 DEVICE_CLASS_MOVING = 'moving'
+
+DEVICE_CLASS_EMPTY = ''
+DEVICE_CLASS_BATTERY = 'battery'
+DEVICE_CLASS_CURRENT = 'current'
+DEVICE_CLASS_ENERGY = 'energy'
+DEVICE_CLASS_HUMIDITY = 'humidity'
+DEVICE_CLASS_ILLUMINANCE = 'illuminance'
+DEVICE_CLASS_SIGNAL_STRENGTH = 'signal_strength'
+DEVICE_CLASS_TEMPERATURE = 'temperature'
+DEVICE_CLASS_POWER = 'power'
+DEVICE_CLASS_POWER_FACTOR = 'power_factor'
+DEVICE_CLASS_PRESSURE = 'pressure'
+DEVICE_CLASS_TIMESTAMP = 'timestamp'
+DEVICE_CLASS_VOLTAGE = 'voltage'

tests/component_tests/sensor/test_sensor.py

@@ -0,0 +1,14 @@
+""" Tests for the sensor component """
+
+
+def test_sensor_device_class_set(generate_main):
+    """
+    When the device_class of sensor is set in the yaml file, it should be registered in main
+    """
+    # Given
+
+    # When
+    main_cpp = generate_main("tests/component_tests/sensor/test_sensor.yaml")
+
+    # Then
+    assert "s_1->set_device_class(\"voltage\");" in main_cpp

tests/component_tests/sensor/test_sensor.yaml

@@ -0,0 +1,12 @@
+esphome:
+  name: test
+  platform: ESP8266
+  board: d1_mini_lite
+
+sensor:
+  - platform: adc
+    pin: A0
+    id: s_1
+    name: "test s1"
+    update_interval: 60s
+    device_class: "voltage"