Adjust type hints for WaterHeaterEntityFeature (home-assistant#82270)

xxv · Nov 17, 2022 · 9d607c8 · 9d607c8
1 parent ced16da
commit 9d607c8
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 26 deletions.
diff --git a/homeassistant/components/atag/water_heater.py b/homeassistant/components/atag/water_heater.py
@@ -30,7 +30,6 @@ class AtagWaterHeater(AtagEntity, WaterHeaterEntity):
     """Representation of an ATAG water heater."""
 
     _attr_operation_list = OPERATION_LIST
-    _attr_supported_features = 0
     _attr_temperature_unit = TEMP_CELSIUS
 
     @property

diff --git a/homeassistant/components/econet/water_heater.py b/homeassistant/components/econet/water_heater.py
@@ -106,7 +106,7 @@ def operation_list(self):
         return op_list
 
     @property
-    def supported_features(self):
+    def supported_features(self) -> WaterHeaterEntityFeature:
         """Return the list of supported features."""
         if self.water_heater.modes:
             if self.water_heater.supports_away:

diff --git a/homeassistant/components/incomfort/water_heater.py b/homeassistant/components/incomfort/water_heater.py
@@ -88,11 +88,6 @@ def temperature_unit(self) -> str:
         """Return the unit of measurement."""
         return TEMP_CELSIUS
 
-    @property
-    def supported_features(self) -> int:
-        """Return the list of supported features."""
-        return 0
-
     @property
     def current_operation(self) -> str:
         """Return the current operation mode."""

diff --git a/homeassistant/components/water_heater/__init__.py b/homeassistant/components/water_heater/__init__.py
@@ -167,7 +167,7 @@ class WaterHeaterEntity(Entity):
     _attr_operation_list: list[str] | None = None
     _attr_precision: float
     _attr_state: None = None
-    _attr_supported_features: int
+    _attr_supported_features: WaterHeaterEntityFeature | int = 0
     _attr_target_temperature_high: float | None = None
     _attr_target_temperature_low: float | None = None
     _attr_target_temperature: float | None = None
@@ -341,6 +341,11 @@ def max_temp(self) -> float:
             DEFAULT_MAX_TEMP, TEMP_FAHRENHEIT, self.temperature_unit
         )
 
+    @property
+    def supported_features(self) -> WaterHeaterEntityFeature | int:
+        """Return the list of supported features."""
+        return self._attr_supported_features
+
 
 async def async_service_away_mode(
     entity: WaterHeaterEntity, service: ServiceCall

diff --git a/pylint/plugins/hass_enforce_type_hints.py b/pylint/plugins/hass_enforce_type_hints.py
@@ -2510,24 +2510,36 @@ class ClassTypeHintMatch:
             base_class="WaterHeaterEntity",
             matches=[
                 TypeHintMatch(
-                    function_name="precision",
-                    return_type="float",
+                    function_name="current_operation",
+                    return_type=["str", None],
                 ),
                 TypeHintMatch(
-                    function_name="temperature_unit",
-                    return_type="str",
+                    function_name="current_temperature",
+                    return_type=["float", None],
                 ),
                 TypeHintMatch(
-                    function_name="current_operation",
-                    return_type=["str", None],
+                    function_name="is_away_mode_on",
+                    return_type=["bool", None],
+                ),
+                TypeHintMatch(
+                    function_name="max_temp",
+                    return_type="float",
+                ),
+                TypeHintMatch(
+                    function_name="min_temp",
+                    return_type="float",
                 ),
                 TypeHintMatch(
                     function_name="operation_list",
                     return_type=["list[str]", None],
                 ),
                 TypeHintMatch(
-                    function_name="current_temperature",
-                    return_type=["float", None],
+                    function_name="precision",
+                    return_type="float",
+                ),
+                TypeHintMatch(
+                    function_name="supported_features",
+                    return_type=["WaterHeaterEntityFeature", "int"],
                 ),
                 TypeHintMatch(
                     function_name="target_temperature",
@@ -2542,8 +2554,8 @@ class ClassTypeHintMatch:
                     return_type=["float", None],
                 ),
                 TypeHintMatch(
-                    function_name="is_away_mode_on",
-                    return_type=["bool", None],
+                    function_name="temperature_unit",
+                    return_type="str",
                 ),
                 TypeHintMatch(
                     function_name="set_temperature",
@@ -2567,14 +2579,6 @@ class ClassTypeHintMatch:
                     return_type=None,
                     has_async_counterpart=True,
                 ),
-                TypeHintMatch(
-                    function_name="min_temp",
-                    return_type="float",
-                ),
-                TypeHintMatch(
-                    function_name="max_temp",
-                    return_type="float",
-                ),
             ],
         ),
     ],