From fb5a67fb1fb1bd701d4b4f0c583afcbb0aff1ec2 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Thu, 18 Aug 2022 19:22:08 +0200 Subject: [PATCH] Add vacuum checks to pylint plugin (#76560) --- pylint/plugins/hass_enforce_type_hints.py | 136 +++++++++++++++++++++- tests/pylint/test_enforce_type_hints.py | 46 +++++++- 2 files changed, 179 insertions(+), 3 deletions(-) diff --git a/pylint/plugins/hass_enforce_type_hints.py b/pylint/plugins/hass_enforce_type_hints.py index 19a69f8808f..4eedca487ab 100644 --- a/pylint/plugins/hass_enforce_type_hints.py +++ b/pylint/plugins/hass_enforce_type_hints.py @@ -55,11 +55,12 @@ class ClassTypeHintMatch: matches: list[TypeHintMatch] +_INNER_MATCH = r"((?:[\w\| ]+)|(?:\.{3})|(?:\w+\[.+\]))" _TYPE_HINT_MATCHERS: dict[str, re.Pattern[str]] = { # a_or_b matches items such as "DiscoveryInfoType | None" - "a_or_b": re.compile(r"^(\w+) \| (\w+)$"), + # or "dict | list | None" + "a_or_b": re.compile(rf"^(.+) \| {_INNER_MATCH}$"), } -_INNER_MATCH = r"((?:[\w\| ]+)|(?:\.{3})|(?:\w+\[.+\]))" _INNER_MATCH_POSSIBILITIES = [i + 1 for i in range(5)] _TYPE_HINT_MATCHERS.update( { @@ -2118,6 +2119,137 @@ _INHERITANCE_MATCH: dict[str, list[ClassTypeHintMatch]] = { ], ), ], + "vacuum": [ + ClassTypeHintMatch( + base_class="Entity", + matches=_ENTITY_MATCH, + ), + ClassTypeHintMatch( + base_class="ToggleEntity", + matches=_TOGGLE_ENTITY_MATCH, + ), + ClassTypeHintMatch( + base_class="_BaseVacuum", + matches=[ + TypeHintMatch( + function_name="battery_level", + return_type=["int", None], + ), + TypeHintMatch( + function_name="battery_icon", + return_type="str", + ), + TypeHintMatch( + function_name="fan_speed", + return_type=["str", None], + ), + TypeHintMatch( + function_name="fan_speed_list", + return_type="list[str]", + ), + TypeHintMatch( + function_name="stop", + kwargs_type="Any", + return_type=None, + has_async_counterpart=True, + ), + TypeHintMatch( + function_name="return_to_base", + kwargs_type="Any", + return_type=None, + has_async_counterpart=True, + ), + TypeHintMatch( + function_name="clean_spot", + kwargs_type="Any", + return_type=None, + has_async_counterpart=True, + ), + TypeHintMatch( + function_name="locate", + kwargs_type="Any", + return_type=None, + has_async_counterpart=True, + ), + TypeHintMatch( + function_name="set_fan_speed", + named_arg_types={ + "fan_speed": "str", + }, + kwargs_type="Any", + return_type=None, + has_async_counterpart=True, + ), + TypeHintMatch( + function_name="send_command", + named_arg_types={ + "command": "str", + "params": "dict[str, Any] | list[Any] | None", + }, + kwargs_type="Any", + return_type=None, + has_async_counterpart=True, + ), + ], + ), + ClassTypeHintMatch( + base_class="VacuumEntity", + matches=[ + TypeHintMatch( + function_name="status", + return_type=["str", None], + ), + TypeHintMatch( + function_name="start_pause", + kwargs_type="Any", + return_type=None, + has_async_counterpart=True, + ), + TypeHintMatch( + function_name="async_pause", + return_type=None, + ), + TypeHintMatch( + function_name="async_start", + return_type=None, + ), + ], + ), + ClassTypeHintMatch( + base_class="StateVacuumEntity", + matches=[ + TypeHintMatch( + function_name="state", + return_type=["str", None], + ), + TypeHintMatch( + function_name="start", + return_type=None, + has_async_counterpart=True, + ), + TypeHintMatch( + function_name="pause", + return_type=None, + has_async_counterpart=True, + ), + TypeHintMatch( + function_name="async_turn_on", + kwargs_type="Any", + return_type=None, + ), + TypeHintMatch( + function_name="async_turn_off", + kwargs_type="Any", + return_type=None, + ), + TypeHintMatch( + function_name="async_toggle", + kwargs_type="Any", + return_type=None, + ), + ], + ), + ], "water_heater": [ ClassTypeHintMatch( base_class="Entity", diff --git a/tests/pylint/test_enforce_type_hints.py b/tests/pylint/test_enforce_type_hints.py index 1381ed34a7b..ebea738edc4 100644 --- a/tests/pylint/test_enforce_type_hints.py +++ b/tests/pylint/test_enforce_type_hints.py @@ -73,7 +73,12 @@ def test_regex_x_of_y_i( @pytest.mark.parametrize( ("string", "expected_a", "expected_b"), - [("DiscoveryInfoType | None", "DiscoveryInfoType", "None")], + [ + ("DiscoveryInfoType | None", "DiscoveryInfoType", "None"), + ("dict | list | None", "dict | list", "None"), + ("dict[str, Any] | list[Any] | None", "dict[str, Any] | list[Any]", "None"), + ("dict[str, Any] | list[Any]", "dict[str, Any]", "list[Any]"), + ], ) def test_regex_a_or_b( hass_enforce_type_hints: ModuleType, string: str, expected_a: str, expected_b: str @@ -967,3 +972,42 @@ def test_number_entity(linter: UnittestLinter, type_hint_checker: BaseChecker) - with assert_no_messages(linter): type_hint_checker.visit_classdef(class_node) + + +def test_vacuum_entity(linter: UnittestLinter, type_hint_checker: BaseChecker) -> None: + """Ensure valid hints are accepted for vacuum entity.""" + # Set bypass option + type_hint_checker.config.ignore_missing_annotations = False + + # Ensure that `dict | list | None` is valid for params + class_node = astroid.extract_node( + """ + class Entity(): + pass + + class ToggleEntity(Entity): + pass + + class _BaseVacuum(Entity): + pass + + class VacuumEntity(_BaseVacuum, ToggleEntity): + pass + + class MyVacuum( #@ + VacuumEntity + ): + def send_command( + self, + command: str, + params: dict[str, Any] | list[Any] | None = None, + **kwargs: Any, + ) -> None: + pass + """, + "homeassistant.components.pylint_test.vacuum", + ) + type_hint_checker.visit_module(class_node.parent) + + with assert_no_messages(linter): + type_hint_checker.visit_classdef(class_node)