Add device class support to Switch (#3012)

Co-authored-by: Oxan van Leeuwen <oxan@oxanvanleeuwen.nl>

esphome/components/api/api.proto

@@ -529,6 +529,7 @@ message ListEntitiesSwitchResponse {
   bool assumed_state = 6;
   bool disabled_by_default = 7;
   EntityCategory entity_category = 8;
+  string device_class = 9;
 }
 message SwitchStateResponse {
   option (id) = 26;

esphome/components/api/api_connection.cpp

@@ -462,6 +462,7 @@ bool APIConnection::send_switch_info(switch_::Switch *a_switch) {
   msg.assumed_state = a_switch->assumed_state();
   msg.disabled_by_default = a_switch->is_disabled_by_default();
   msg.entity_category = static_cast<enums::EntityCategory>(a_switch->get_entity_category());
+  msg.device_class = a_switch->get_device_class();
   return this->send_list_entities_switch_response(msg);
 }
 void APIConnection::switch_command(const SwitchCommandRequest &msg) {

esphome/components/api/api_pb2.cpp

@@ -2177,6 +2177,10 @@ bool ListEntitiesSwitchResponse::decode_length(uint32_t field_id, ProtoLengthDel
       this->icon = value.as_string();
       return true;
     }
+    case 9: {
+      this->device_class = value.as_string();
+      return true;
+    }
     default:
       return false;
   }
@@ -2200,6 +2204,7 @@ void ListEntitiesSwitchResponse::encode(ProtoWriteBuffer buffer) const {
   buffer.encode_bool(6, this->assumed_state);
   buffer.encode_bool(7, this->disabled_by_default);
   buffer.encode_enum<enums::EntityCategory>(8, this->entity_category);
+  buffer.encode_string(9, this->device_class);
 }
 #ifdef HAS_PROTO_MESSAGE_DUMP
 void ListEntitiesSwitchResponse::dump_to(std::string &out) const {
@@ -2237,6 +2242,10 @@ void ListEntitiesSwitchResponse::dump_to(std::string &out) const {
   out.append("  entity_category: ");
   out.append(proto_enum_to_string<enums::EntityCategory>(this->entity_category));
   out.append("\n");
+
+  out.append("  device_class: ");
+  out.append("'").append(this->device_class).append("'");
+  out.append("\n");
   out.append("}");
 }
 #endif

esphome/components/api/api_pb2.h

@@ -580,6 +580,7 @@ class ListEntitiesSwitchResponse : public ProtoMessage {
   bool assumed_state{false};
   bool disabled_by_default{false};
   enums::EntityCategory entity_category{};
+  std::string device_class{};
   void encode(ProtoWriteBuffer buffer) const override;
 #ifdef HAS_PROTO_MESSAGE_DUMP
   void dump_to(std::string &out) const override;

esphome/components/switch/__init__.py

@@ -4,18 +4,27 @@ from esphome import automation
 from esphome.automation import Condition, maybe_simple_id
 from esphome.components import mqtt
 from esphome.const import (
+    CONF_DEVICE_CLASS,
     CONF_ID,
     CONF_INVERTED,
+    CONF_MQTT_ID,
     CONF_ON_TURN_OFF,
     CONF_ON_TURN_ON,
     CONF_TRIGGER_ID,
-    CONF_MQTT_ID,
+    DEVICE_CLASS_EMPTY,
+    DEVICE_CLASS_OUTLET,
+    DEVICE_CLASS_SWITCH,
 )
 from esphome.core import CORE, coroutine_with_priority
 from esphome.cpp_helpers import setup_entity
 
 CODEOWNERS = ["@esphome/core"]
 IS_PLATFORM_COMPONENT = True
+DEVICE_CLASSES = [
+    DEVICE_CLASS_EMPTY,
+    DEVICE_CLASS_OUTLET,
+    DEVICE_CLASS_SWITCH,
+]
 
 switch_ns = cg.esphome_ns.namespace("switch_")
 Switch = switch_ns.class_("Switch", cg.EntityBase)
@@ -51,6 +60,7 @@ SWITCH_SCHEMA = cv.ENTITY_BASE_SCHEMA.extend(cv.MQTT_COMMAND_COMPONENT_SCHEMA).e
                 cv.GenerateID(CONF_TRIGGER_ID): cv.declare_id(SwitchTurnOffTrigger),
             }
         ),
+        cv.Optional(CONF_DEVICE_CLASS): cv.one_of(*DEVICE_CLASSES, lower=True),
     }
 )
 
@@ -71,6 +81,9 @@ async def setup_switch_core_(var, config):
         mqtt_ = cg.new_Pvariable(config[CONF_MQTT_ID], var)
         await mqtt.register_mqtt_component(mqtt_, config)
 
+    if CONF_DEVICE_CLASS in config:
+        cg.add(var.set_device_class(config[CONF_DEVICE_CLASS]))
+
 
 async def register_switch(var, config):
     if not CORE.has_id(config[CONF_ID]):

esphome/components/switch/switch.cpp

@@ -46,5 +46,12 @@ void Switch::set_inverted(bool inverted) { this->inverted_ = inverted; }
 uint32_t Switch::hash_base() { return 3129890955UL; }
 bool Switch::is_inverted() const { return this->inverted_; }
 
+std::string Switch::get_device_class() {
+  if (this->device_class_.has_value())
+    return *this->device_class_;
+  return "";
+}
+void Switch::set_device_class(const std::string &device_class) { this->device_class_ = device_class; }
+
 }  // namespace switch_
 }  // namespace esphome

esphome/components/switch/switch.h

@@ -20,6 +20,9 @@ namespace switch_ {
     if ((obj)->is_inverted()) { \
       ESP_LOGCONFIG(TAG, "%s  Inverted: YES", prefix); \
     } \
+    if (!(obj)->get_device_class().empty()) { \
+      ESP_LOGCONFIG(TAG, "%s  Device Class: '%s'", prefix, (obj)->get_device_class().c_str()); \
+    } \
   }
 
 /** Base class for all switches.
@@ -88,6 +91,11 @@ class Switch : public EntityBase {
 
   bool is_inverted() const;
 
+  /// Get the device class for this switch.
+  std::string get_device_class();
+  /// Set the Home Assistant device class for this switch.
+  void set_device_class(const std::string &device_class);
+
  protected:
   /** Write the given state to hardware. You should implement this
    * abstract method if you want to create your own switch.
@@ -105,6 +113,7 @@ class Switch : public EntityBase {
   bool inverted_{false};
   Deduplicator<bool> publish_dedup_;
   ESPPreferenceObject rtc_;
+  optional<std::string> device_class_;
 };
 
 }  // namespace switch_

esphome/const.py

@@ -920,6 +920,9 @@ DEVICE_CLASS_VOLTAGE = "voltage"
 DEVICE_CLASS_UPDATE = "update"
 # device classes of button component
 DEVICE_CLASS_RESTART = "restart"
+# device classes of switch component
+DEVICE_CLASS_OUTLET = "outlet"
+DEVICE_CLASS_SWITCH = "switch"
 
 
 # state classes

tests/test1.yaml

@@ -2037,6 +2037,10 @@ switch:
   - platform: template
     id: ble1_status
     optimistic: true
+  - platform: template
+    id: outlet_switch
+    optimistic: true
+    device_class: outlet
 
 fan:
   - platform: binary