Add WaterHeaterFeature class for ESPHome (#1493)

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: J. Nick Koston <nick@home-assistant.io>

Author: 4 people

aioesphomeapi/api.proto

@@ -1155,9 +1155,11 @@ enum WaterHeaterCommandHasField {
   WATER_HEATER_COMMAND_HAS_NONE = 0;
   WATER_HEATER_COMMAND_HAS_MODE = 1;
   WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE = 2;
-  WATER_HEATER_COMMAND_HAS_STATE = 4;
+  WATER_HEATER_COMMAND_HAS_STATE = 4 [deprecated=true];
   WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_LOW = 8;
   WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_HIGH = 16;
+  WATER_HEATER_COMMAND_HAS_ON_STATE = 32;
+  WATER_HEATER_COMMAND_HAS_AWAY_STATE = 64;
 }
 
 message WaterHeaterCommandRequest {

aioesphomeapi/api_pb2.py

@@ -1405,14 +1405,19 @@ def water_heater_command(
             req.has_fields |= WaterHeaterCommandField.TARGET_TEMPERATURE
             req.target_temperature = target_temperature
 
-        if away is not None or on is not None:
-            req.has_fields |= WaterHeaterCommandField.STATE
-            state = WaterHeaterStateFlag(0)
+        # Use the new granular AWAY_STATE/ON_STATE fields instead of the
+        # deprecated combined STATE field. Water heater support only shipped
+        # in HA 2026.2.0/ESPHome 2026.1.0 so early adopters can update
+        # their devices; no need for legacy STATE backward compat.
+        if away is not None:
+            req.has_fields |= WaterHeaterCommandField.AWAY_STATE
             if away:
-                state |= WaterHeaterStateFlag.AWAY
+                req.state |= WaterHeaterStateFlag.AWAY
+
+        if on is not None:
+            req.has_fields |= WaterHeaterCommandField.ON_STATE
             if on:
-                state |= WaterHeaterStateFlag.ON
-            req.state = state
+                req.state |= WaterHeaterStateFlag.ON
 
         if target_temperature_low is not None:
             req.has_fields |= WaterHeaterCommandField.TARGET_TEMPERATURE_LOW

aioesphomeapi/client.py

@@ -1109,6 +1109,16 @@ class WaterHeaterMode(APIIntEnum):
     GAS = 6
 
 
+class WaterHeaterFeature(enum.IntFlag):
+    """ESPHome water heater feature flags (WaterHeaterFeature)."""
+
+    SUPPORTS_CURRENT_TEMPERATURE = 1 << 0
+    SUPPORTS_TARGET_TEMPERATURE = 1 << 1
+    SUPPORTS_OPERATION_MODE = 1 << 2
+    SUPPORTS_AWAY_MODE = 1 << 3
+    SUPPORTS_ON_OFF = 1 << 4
+
+
 class WaterHeaterCommandField(enum.IntFlag):
     """Bitmask for has_fields in WaterHeaterCommandRequest."""
 
@@ -1117,6 +1127,8 @@ class WaterHeaterCommandField(enum.IntFlag):
     STATE = 1 << 2
     TARGET_TEMPERATURE_LOW = 1 << 3
     TARGET_TEMPERATURE_HIGH = 1 << 4
+    ON_STATE = 1 << 5
+    AWAY_STATE = 1 << 6
 
 
 class WaterHeaterStateFlag(enum.IntFlag):

aioesphomeapi/model.py

@@ -760,7 +760,7 @@ async def test_switch_command(
             {
                 "key": 1,
                 "device_id": 0,
-                "has_fields": WaterHeaterCommandField.STATE,
+                "has_fields": WaterHeaterCommandField.AWAY_STATE,
                 "state": WaterHeaterStateFlag.AWAY,
             },
         ),
@@ -769,7 +769,7 @@ async def test_switch_command(
             {
                 "key": 1,
                 "device_id": 0,
-                "has_fields": WaterHeaterCommandField.STATE,
+                "has_fields": WaterHeaterCommandField.ON_STATE,
                 "state": WaterHeaterStateFlag.ON,
             },
         ),
@@ -778,7 +778,8 @@ async def test_switch_command(
             {
                 "key": 1,
                 "device_id": 0,
-                "has_fields": WaterHeaterCommandField.STATE,
+                "has_fields": WaterHeaterCommandField.AWAY_STATE
+                | WaterHeaterCommandField.ON_STATE,
                 "state": WaterHeaterStateFlag.AWAY | WaterHeaterStateFlag.ON,
             },
         ),
@@ -787,7 +788,8 @@ async def test_switch_command(
             {
                 "key": 1,
                 "device_id": 0,
-                "has_fields": WaterHeaterCommandField.STATE,
+                "has_fields": WaterHeaterCommandField.AWAY_STATE
+                | WaterHeaterCommandField.ON_STATE,
                 "state": WaterHeaterStateFlag(0),
             },
         ),