Enforce VacuumEntityFeature (#82466)

home-assistant · Nov 22, 2022 · bf3c6e5 · bf3c6e5
1 parent 4134d72
commit bf3c6e5
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 12 deletions.
diff --git a/homeassistant/components/demo/vacuum.py b/homeassistant/components/demo/vacuum.py
@@ -95,7 +95,7 @@ async def async_setup_platform(
             DemoVacuum(DEMO_VACUUM_MOST, SUPPORT_MOST_SERVICES),
             DemoVacuum(DEMO_VACUUM_BASIC, SUPPORT_BASIC_SERVICES),
             DemoVacuum(DEMO_VACUUM_MINIMAL, SUPPORT_MINIMAL_SERVICES),
-            DemoVacuum(DEMO_VACUUM_NONE, 0),
+            DemoVacuum(DEMO_VACUUM_NONE, VacuumEntityFeature(0)),
             StateDemoVacuum(DEMO_VACUUM_STATE),
         ]
     )
@@ -106,9 +106,7 @@ class DemoVacuum(VacuumEntity):
 
     _attr_should_poll = False
 
-    def __init__(
-        self, name: str, supported_features: VacuumEntityFeature | int
-    ) -> None:
+    def __init__(self, name: str, supported_features: VacuumEntityFeature) -> None:
         """Initialize the vacuum."""
         self._attr_name = name
         self._attr_supported_features = supported_features

diff --git a/homeassistant/components/mqtt/vacuum/schema.py b/homeassistant/components/mqtt/vacuum/schema.py
@@ -20,7 +20,7 @@
 
 
 def services_to_strings(
-    services: VacuumEntityFeature | int,
+    services: VacuumEntityFeature,
     service_to_string: dict[VacuumEntityFeature, str],
 ) -> list[str]:
     """Convert SUPPORT_* service bitmask to list of service strings."""
@@ -33,9 +33,9 @@ def services_to_strings(
 
 def strings_to_services(
     strings: list[str], string_to_service: dict[str, VacuumEntityFeature]
-) -> VacuumEntityFeature | int:
+) -> VacuumEntityFeature:
     """Convert service strings to SUPPORT_* service bitmask."""
-    services: VacuumEntityFeature | int = 0
+    services = VacuumEntityFeature(0)
     for string in strings:
         services |= string_to_service[string]
     return services
diff --git a/homeassistant/components/mqtt/vacuum/schema_state.py b/homeassistant/components/mqtt/vacuum/schema_state.py
@@ -60,15 +60,15 @@
 STRING_TO_SERVICE = {v: k for k, v in SERVICE_TO_STRING.items()}
 
 
-DEFAULT_SERVICES: VacuumEntityFeature | int = (
+DEFAULT_SERVICES = (
     VacuumEntityFeature.START
     | VacuumEntityFeature.STOP
     | VacuumEntityFeature.RETURN_HOME
     | VacuumEntityFeature.STATUS
     | VacuumEntityFeature.BATTERY
     | VacuumEntityFeature.CLEAN_SPOT
 )
-ALL_SERVICES: VacuumEntityFeature | int = (
+ALL_SERVICES = (
     DEFAULT_SERVICES
     | VacuumEntityFeature.PAUSE
     | VacuumEntityFeature.LOCATE

diff --git a/homeassistant/components/vacuum/__init__.py b/homeassistant/components/vacuum/__init__.py
@@ -180,10 +180,10 @@ class _BaseVacuum(Entity):
     _attr_battery_level: int | None = None
     _attr_fan_speed: str | None = None
     _attr_fan_speed_list: list[str]
-    _attr_supported_features: VacuumEntityFeature | int = 0
+    _attr_supported_features: VacuumEntityFeature = VacuumEntityFeature(0)
 
     @property
-    def supported_features(self) -> VacuumEntityFeature | int:
+    def supported_features(self) -> VacuumEntityFeature:
         """Flag vacuum cleaner features that are supported."""
         return self._attr_supported_features
 

diff --git a/pylint/plugins/hass_enforce_type_hints.py b/pylint/plugins/hass_enforce_type_hints.py
@@ -2396,7 +2396,7 @@ class ClassTypeHintMatch:
                 ),
                 TypeHintMatch(
                     function_name="supported_features",
-                    return_type=["VacuumEntityFeature", "int"],
+                    return_type="VacuumEntityFeature",
                 ),
                 TypeHintMatch(
                     function_name="stop",